#!/usr/bin/env python

# unit tests for canon module

import unittest

import testbase

import canon
import absyn
import ir
import fsymbol
import copy
from fracttypes import *

class CanonTest(testbase.TestBase):
    def setUp(self):
        self.fakeNode = absyn.Empty(0)
        self.canon = canon.T(fsymbol.T())
    def tearDown(self):
        pass

    # convenience methods to make quick trees for testing
    def eseq(self,stms, exp):
        return ir.ESeq(stms, exp, self.fakeNode, Int)
    def seq(self,stms):
        return ir.Seq(stms,self.fakeNode)
    def var(self,name="a"):
        return ir.Var(name,self.fakeNode, Int)
    def const(self,value=0):
        return ir.Const(value, self.fakeNode, Int)
    def binop(self,stms,op="+"):
        return ir.Binop(op,stms,self.fakeNode, Int)
    def move(self,dest,exp):
        return ir.Move(dest, exp, self.fakeNode, Int)
    def cjump(self,e1,e2,trueDest="trueDest",falseDest="falseDest"):
        return ir.CJump(">", e1, e2, trueDest, falseDest, self.fakeNode)
    def jump(self,dest):
        return ir.Jump(dest,self.fakeNode)
    def cast(self, e, type):
        return ir.Cast(e,self.fakeNode, type)
    def label(self,name):
        return ir.Label(name,self.fakeNode)
    
    def testEmptyTree(self):
        self.assertEqual(self.canon.linearize(None),None)
                                     
    def testBinop(self):
        # binop with no eseqs
        tree = self.binop([self.var(), self.const()])
        ltree = self.canon.linearize(tree)
        self.assertTreesEqual("",tree, ltree)
        self.assertESeqsNotNested(ltree,1)

    def testLHBinop(self):
        # left-hand eseq
        tree = self.binop([self.eseq([self.move(self.var(),self.const())],
                                      self.var("b")),
                           self.const()])

        ltree = self.canon.linearize(tree)
        self.failUnless(isinstance(ltree,ir.ESeq) and \
                        isinstance(ltree.children[0],ir.Move) and \
                        isinstance(ltree.children[1],ir.Binop) and \
                        isinstance(ltree.children[1].children[0],ir.Var))
        self.assertESeqsNotNested(ltree,1)

        # nested left-hand eseq
        tree = self.binop([self.eseq([self.move(self.var(),self.const())],
                                      self.var("b")),
                           self.const()])

        tree = self.binop([tree,self.const()])

        ltree = self.canon.linearize(tree)
        self.assertESeqsNotNested(ltree,1)

    def testRHBinop(self):
        # right-hand eseq
        tree = self.binop([self.var("a"),
                           self.eseq([self.move(self.var("b"),self.const())],
                                                self.var("b"))])
        ltree = self.canon.linearize(tree)
        self.assertESeqsNotNested(ltree,1)
        self.failUnless(isinstance(ltree.children[0].children[0], ir.Var) and \
                        ltree.children[0].children[0].name == \
                        ltree.children[2].children[0].name)

        # commuting right-hand eseq
        tree = self.binop([self.const(4),
                           self.eseq([self.move(self.var("b"),self.const())],
                                                self.var("b"))])
        ltree = self.canon.linearize(tree)
        self.assertESeqsNotNested(ltree,1)
        self.failUnless(isinstance(ltree.children[1].children[0],ir.Const))

    def testNestRHBinop(self):
        # nested right-hand eseq
        tree = self.binop([self.var("a"),
                           self.eseq([self.move(self.var("b"),self.const())],
                                                self.var("b"))])
        tree = self.binop([self.const(),tree])

        # what we ought to produce
        exptree = self.eseq([self.move(self.var("t__0"), self.var()),
                             self.move(self.var("b"),self.const())],
                            self.binop([self.const(),
                                        self.binop([self.var("t__0"),
                                                    self.var("b")])]))
        ltree = self.canon.linearize(tree)
        self.assertESeqsNotNested(ltree,1)
        self.assertTreesEqual("",ltree,exptree)

    def testBothSidesBinop(self):
        tree = self.binop([self.var("a"),
                           self.eseq([self.move(self.var("b"),self.const())],
                                                self.var("b"))])
        tree2 = copy.deepcopy(tree)
        tree = self.binop([tree, tree2])
        ltree = self.canon.linearize(tree)
        self.assertESeqsNotNested(ltree,1)
        
    def testESeq(self):
        tree = self.eseq([self.var("a")], self.var("b"))
        ltree = self.canon.linearize(tree)
        self.assertTreesEqual("",tree,ltree)
        self.assertESeqsNotNested(ltree,1)
        
        tree = self.eseq([self.eseq([self.var("a")], self.var("b"))], self.var("c"))
        ltree = self.canon.linearize(tree)
        self.assertESeqsNotNested(ltree,1)
        
        tree2 = self.eseq([self.var("a")], self.eseq([self.var("b")], self.var("c")))
        ltree2 = self.canon.linearize(tree2)
        self.assertESeqsNotNested(ltree2,1)
        self.assertTreesEqual("",ltree, ltree2)

    def testSeq(self):
        tree = self.seq([self.var("a"),self.var("b")])
        ltree = self.canon.linearize(tree)
        self.assertTreesEqual("",tree,ltree)
        self.assertESeqsNotNested(ltree,1)

        tree = self.seq([self.seq([self.var("a"), self.var("b")]), self.var("c")])
        ltree = self.canon.linearize(tree)
        self.assertESeqsNotNested(ltree,1)
        
        tree2 = self.seq([self.var("a"), self.seq([self.var("b"), self.var("c")])])
        ltree2 = self.canon.linearize(tree2)
        self.assertESeqsNotNested(ltree2,1)
        self.assertTreesEqual("",ltree, ltree2)
        
    def testCJump(self):
        tree = self.cjump(self.var(), self.var())
        ltree = self.canon.linearize(tree)
        self.assertESeqsNotNested(ltree,1)
        self.assertTreesEqual("",tree, ltree)

    def testCJumpWithLHESeq(self):
        tree = self.cjump(self.eseq([self.move(self.var("a"),self.const())],
                                    self.var("b")),
                          self.var("c"))
        ltree = self.canon.linearize(tree)
        self.assertESeqsNotNested(ltree,1)

    def testCJumpWithRHESeq(self):
        # commutes
        tree = self.cjump(self.const(),
                          self.eseq([self.move(self.var("a"),self.const())],
                                    self.var("b")))
        ltree = self.canon.linearize(tree)
        self.assertESeqsNotNested(ltree,1)
        
        # doesn't commute
        tree = self.cjump(self.var("c"),
                          self.eseq([self.move(self.var("a"),self.const())],
                                    self.var("b")))
        ltree = self.canon.linearize(tree)
        self.assertESeqsNotNested(ltree,1)

    def testBothSidesCJump(self):
        tree = self.cjump(self.eseq([self.move(self.var("a"),self.const())],
                                    self.var("b")),
                          self.eseq([self.move(self.var("a"),self.const())],
                                    self.var("b")))
        ltree = self.canon.linearize(tree)
        self.assertESeqsNotNested(ltree,1)

    def testCast(self):
        tree = self.cast(self.eseq([self.const(4)], self.var()), Float)

        ltree = self.canon.linearize(tree)
        self.assertESeqsNotNested(ltree,1)

    def testBasicBlocks(self):
        # no jumps or labels
        seq = self.seq([self.var(), self.var("b"), self.var("c")])
        blocks = self.canon.basic_blocks(seq, "t__start", "t__end")
        self.assertBlockIsWellFormed(blocks[0],"t__start","t__end")
        trace = self.canon.schedule_trace(blocks,"t__end")
        self.assertValidTrace(trace)

        # starts with a label
        seq = self.seq([self.label("t__1"),self.var()])
        blocks = self.canon.basic_blocks(seq, "t__start", "t__end")
        self.assertBlockIsWellFormed(blocks[0],"t__1","t__end")

        # just a label
        seq = self.seq([self.label("t__1")])
        blocks = self.canon.basic_blocks(seq, "t__start", "t__end")
        self.assertBlockIsWellFormed(blocks[0],"t__1","t__end")

        # just a jump
        seq = self.seq([self.jump("d__1")])
        blocks = self.canon.basic_blocks(seq, "t__start", "t__end")
        self.assertBlockIsWellFormed(blocks[0],"t__start","d__1")

        # empty seq
        seq = self.seq([])
        blocks = self.canon.basic_blocks(seq, "t__start", "t__end")
        self.assertEqual(blocks,[])

        # jump midway
        seq = self.seq([self.var("a"),self.jump("d__1"),self.var("b")])
        blocks = self.canon.basic_blocks(seq, "t__start", "t__end")
        self.assertBlocksAreWellFormed(blocks)
        self.assertEqual(len(blocks),2)

        # cjump midway
        seq = self.seq([self.var("a"),self.cjump(self.var(),self.var()),self.var("b")])
        blocks = self.canon.basic_blocks(seq, "t__start", "t__end")
        self.assertBlocksAreWellFormed(blocks)
        self.assertEqual(len(blocks),2)

        # label midway
        seq = self.seq([self.var("a"), self.label("x"), self.var("b")])
        blocks = self.canon.basic_blocks(seq, "t__start", "t__end")
        self.assertBlocksAreWellFormed(blocks)
        self.assertEqual(len(blocks),2)

        # starts with a jump and ends with label
        seq = self.seq([self.jump("d"),self.var("a"), self.label("x")])
        blocks = self.canon.basic_blocks(seq, "t__start", "t__end")
        self.assertBlocksAreWellFormed(blocks)
        self.assertEqual(len(blocks),3)

        # jump to next stm
        seq = self.seq([self.jump("d"),self.label("d")])        
        blocks = self.canon.basic_blocks(seq, "t__start", "t__end")
        self.assertBlocksAreWellFormed(blocks)
        self.assertEqual(len(blocks),2)        

    def testTraceScheduling(self):
        # cjump followed by falsedest
        seq = self.seq([self.var("a"),
                        self.cjump(self.var(),self.var(),"t__end","t__0"),
                        self.label("t__0"),
                        self.var("b")])
        blocks = self.canon.basic_blocks(seq, "t__start", "t__end")
        trace = self.canon.schedule_trace(blocks,"t__end")
        self.assertValidTrace(trace)

        # cjump followed by truedest
        seq = self.seq([self.var("a"),
                        self.cjump(self.var(),self.var(),"t__0","t__end"),
                        self.label("t__0"),
                        self.var("b")])
        blocks = self.canon.basic_blocks(seq, "t__start", "t__end")
        trace = self.canon.schedule_trace(blocks,"t__end")
        self.assertValidTrace(trace)
        self.assertEqual(trace[2].op,"<=")

        # cjump followed by neither
        seq = self.seq([self.label("start"),
                        self.cjump(self.var(),self.var(),"start","t__end"),
                        self.label("foo"),
                        self.const()])
        blocks = self.canon.basic_blocks(seq, "t__start", "t__end")
        trace = self.canon.schedule_trace(blocks,"t__end")
        self.assertValidTrace(trace)

        # two mingled traces
        seq = self.seq([self.cjump(self.const(), self.const(),"a1","b1"),
                        self.label("a1"), self.var("a1"),self.jump("a2"),
                        self.label("b1"), self.var("b1"),self.jump("b2"),
                        self.label("a2"), self.var("a2"),self.jump("t__end"),
                        self.label("b2"), self.var("b2"),self.jump("t__end")])
        blocks = self.canon.basic_blocks(seq, "t__start", "t__end")
        trace = self.canon.schedule_trace(blocks,"t__end")
        self.assertValidTrace(trace)
        self.failUnless(trace[2].name == "b1" and \
                        trace[3].name == "b1" and \
                        trace[4].name == "b2")
        self.failUnless(trace[7].name == "a1" and \
                        trace[8].name == "a1" and \
                        trace[9].name == "a2")

    def test_canonicalize(self):
        # check overall driver works ok
        seq = self.seq([self.cjump(self.const(), self.const(),"a1","b1"),
                        self.label("a1"), self.var("a1"),self.jump("a2"),
                        self.label("b1"), self.var("b1"),self.jump("b2"),
                        self.label("a2"), self.var("a2"),self.jump("t__end"),
                        self.label("b2"), self.var("b2"),self.jump("t__end")])

        trace = self.canon.canonicalize(seq,"t__start","t__end")
        self.assertValidTrace(trace)
        
    def printAllBlocks(self,blocks):
        for b in blocks:
            for stm in b: print stm.pretty(),
            print

#     def assertValidTrace(self,trace):
#         # must have each cjump followed by false case
#         expecting = None
#         for stm in trace:
#             if expecting != None:
#                 self.failUnless(isinstance(stm,ir.Label))
#                 self.assertEqual(stm.name,expecting)
#                 expecting = None
#             elif isinstance(stm, ir.CJump):
#                 expecting = stm.falseDest
                
                    
#     def assertESeqsNotNested(self,t,parentAllowsESeq):
#         'check that no ESeqs are left below other nodes'
#         if isinstance(t,ir.ESeq):
#             if parentAllowsESeq:
#                 for child in t.children:
#                     self.assertESeqsNotNested(child,0)
#             else:
#                 self.fail("tree not well-formed after linearize")
#         else:
#             for child in t.children:
#                 self.assertESeqsNotNested(child,0)
                
#     def assertTreesEqual(self, t1, t2):
#         self.failUnless(
#             t1.pretty() == t2.pretty(),
#             ("%s, %s should be equivalent" % (t1.pretty(), t2.pretty())))

def suite():
    return unittest.makeSuite(CanonTest,'test')

if __name__ == '__main__':
    unittest.main(defaultTest='suite')