package com.seaofnodes.simple;
import com.seaofnodes.simple.type.*;
import com.seaofnodes.simple.node.*;


import com.seaofnodes.simple.codegen.CodeGen.Phase;
import com.seaofnodes.simple.codegen.CodeGen;
import org.junit.Ignore;
import org.junit.Test;
import static org.junit.Assert.*;
import static org.junit.Assert.fail;

public class Chapter18Test {
    @Test public void testFunctionLocalConstantChains() {
        String src = "val f = { int x -> x ? f(x-1)*305420988+x/305420988 : 1; }; "+
                     "val g = { int x -> x ? g(x-1)/305420988+x*305420988 : 1; }; "+
                     "return f(arg)+g(arg);";
        CodeGen code = new CodeGen(src);
        code.driver(CodeGen.Phase.Schedule,"riscv","SystemV");
        var owners = new java.util.IdentityHashMap<CFGNode,Boolean>();
        code._stop.walk(n -> {
            if( !(n instanceof com.seaofnodes.simple.node.cpus.riscv.AddIRISC add) ||
                !(add.in(1) instanceof com.seaofnodes.simple.node.cpus.riscv.LUI upper) ||
                ((TypeInteger)upper._con).value()+((add._imm12<<20)>>20)!=305420988 )
                return null;
            CFGNode fun = n.cfg0();
            while( fun!=null && !(fun instanceof FunNode) ) fun = fun.idom();
            assertTrue("Constant must belong to a function",fun instanceof FunNode);
            assertNull("Share one constant chain within each function",owners.put(fun,true));
            // Follow all parts of a machine constant, including shared inputs.
            var todo = new java.util.ArrayList<Node>();
            var seen = new java.util.IdentityHashMap<Node,Boolean>();
            todo.add(n);
            for( int j=0; j<todo.size(); j++ ) {
                Node part = todo.get(j);
                if( seen.put(part,true)!=null ) continue;
                CFGNode owner = part.cfg0();
                while( owner!=null && !(owner instanceof FunNode) ) owner = owner.idom();
                assertSame("Every constant-building operation is function-local",fun,owner);
                for( int i=1; i<part.nIns(); i++ ) {
                    Node def = part.in(i);
                    if( def==null || def instanceof CFGNode ) continue;
                    assertTrue("Copies must register their data edges",def._outputs.find(part)!=-1);
                    if( def._type.isConstant() ) todo.add(def);
                }
            }
            return null;
        });
        assertTrue("Exercise constants shared across functions",owners.size()>=2);
    }

    @Test public void testReachableMixedReturns() {
        for( String[] test : new String[][] {
            {"if(arg) return 7; else return 2.5;", "No common type amongst int and f64"},
            {"struct S { int x; }; if(arg) return new S; else return 0;",
             "No common type amongst int and reference"}
        } ) {
            try {
                new CodeGen(test[0]).parse().opto().typeCheck();
                fail("Expected incompatible return types: " + test[0]);
            } catch( Parser.ParseException e ) {
                assertEquals("Cannot return generic scalar",e.getMessage());
            }
        }
    }


    @Test public void testPrintingForwardReferenceScope() throws Exception {
        new CodeGen("return 0;").parse();
        var declared = TypeMemPtr.make(TypeStruct.make("PrintForward",true,true));
        var scope = new ScopeNode();
        scope.define("ptr",declared,false,new ConstantNode(TypeInteger.ZERO),null);
        Parser.TYPES.put("PrintForward",TypeMemPtr.make(TypeStruct.make("PrintForward",false,false)));
        var field = Var.class.getDeclaredField("_type");
        field.setAccessible(true);
        var v = scope.var(0);
        org.junit.Assert.assertSame(declared,field.get(v));
        scope.toString();
        org.junit.Assert.assertSame(declared,field.get(v));
    }


    @Test
    public void testJig() {
        CodeGen code = new CodeGen(
"""
return 0;
""");
        code.driver(Phase.LocalSched);
        assertEquals("return 0;", code.print());
        assertEquals("0", Eval2.eval(code,  2));
    }

    @Test
    public void testPhiParalleAssign() {
        CodeGen code = new CodeGen(
"""
int a = 1;
int b = 2;
while(arg--) {
  int t = a;
  a = b;
  b = t;
}
return a;
""");
        code.parse().opto().typeCheck();
        assertEquals("return Phi(Loop,1,Phi(Loop,2,Phi_a));", code.print());
        assertEquals("1", Eval2.eval(code,  0));
        assertEquals("2", Eval2.eval(code,  1));
        assertEquals("1", Eval2.eval(code,  2));
        assertEquals("2", Eval2.eval(code,  3));
    }


    // ---------------------------------------------------------------
    @Test
    public void testType0() {
        CodeGen code = new CodeGen(
"""
{int -> int}? x2 = null; // null function ptr
return x2;
""");
        code.parse().opto();
        assertEquals("return null;", code.print());
        assertEquals("null", Eval2.eval(code, 0 ) );
    }

    @Test
    public void testFcn0() {
        CodeGen code = new CodeGen(
"""
{int -> int}? sq = { int x ->
    x*x;
};
""");
        code.parse().opto();
        assertEquals("Stop[ return { sq}; return (Parm_x(sq,i64)*x); ]", code.print());
        assertEquals("{ ptr i64 -> i64 #[ 3]}", Eval2.eval(code, 3));
    }

    @Test
    public void testFcn1() {
        CodeGen code = new CodeGen(
"""
var _sq = { int x ->
    x*x;
};
return _sq(arg)+_sq(3);
""");
        code.driver(Phase.LocalSched);
        assertEquals("return ((arg*arg)+9);", code.print());
        assertEquals("13", Eval2.eval(code, 2));
    }

    // Function scope test
    @Test
    public void testFcn2A() {
        CodeGen code = new CodeGen(
"""
int cnt=1;                      // Global scope; exactly one of these
return { -> cnt++; };           // Escape global-scope variable is OK
""");
        code.parse().opto();
        assertEquals("Stop[ return { ptr -> i64 #[ 3]}; return .cnt; ]", code.print());
        assertEquals("{ ptr -> i64 #[ 3]}", Eval2.eval(code, 0));
    }

    // Function scope test
    @Test
    public void testFcn2B() {
        CodeGen code = new CodeGen(
"""
val ctr = { ->
    int cnt; // Function-local var
    return { -> cnt=cnt+1; cnt; };       // Escape cnt out of scope
};
val A = ctr();                  // Make a counter
val B = ctr();                  // Make a counter
A();                            // Increment counter A
return A()*10 + B();            // Incr A and B; result: 2*10+1 == 21
""");
        try { code.parse().opto(); fail(); }
        catch( Exception e ) { assertEquals("Variable 'cnt' is out of function scope and must be a final constant",e.getMessage()); }
    }

    // Function scope test
    @Test
    public void testFcn3() {
        CodeGen code = new CodeGen(
"""
val cnt=2;
return { -> cnt; }();
""");
        code.parse().opto();
        assertEquals("return 2;", code.print());
        assertEquals("2", Eval2.eval(code, 0));
    }

    // Function variables
    @Test
    public void testFcn4() {
        CodeGen code = new CodeGen(
"""
var fcn = arg ? { int x -> x*x; } : { int x -> x+x; };
return fcn(3);
""");
        code.parse().opto();
        assertEquals("Stop[ return #2; return (Parm_x($fun3,i64,3)*x); return (Parm_x($fun4,i64,3)<<1); ]", code.print());
        assertEquals("6", Eval2.eval(code, 0));
        assertEquals("9", Eval2.eval(code, 1));
    }

    // Recursive factorial test
    @Test
    public void testFcn5() {
        CodeGen code = new CodeGen("val _fact = { int x -> x <= 1 ? 1 : x*_fact(x-1); }; return _fact(arg);");
        code.parse().opto().typeCheck();
        String ir = code.print();
        assertTrue(ir.equals("Stop[ return #2; return Phi(Region,1,(Parm_x(_fact,arg,(x-1))*#2)); ]") ||
                   ir.equals("Stop[ return Phi(Region,1,(arg*#2)); return Phi(Region,1,(Parm_x(_fact,(arg-1),(x-1))*#2)); ]") ||
                   ir.equals("Stop[ return Phi(Region,1,(arg*Phi(Region,1,((arg-1)*#2)))); return Phi(Region,1,(Parm_x(_fact,(Sub-1),(x-1))*#2)); ]"));
        assertEquals( "1", Eval2.eval(code, 0));
        assertEquals( "1", Eval2.eval(code, 1));
        assertEquals( "2", Eval2.eval(code, 2));
        assertEquals( "6", Eval2.eval(code, 3));
        assertEquals("24", Eval2.eval(code, 4));
    }

    // The multiply initially sees one integer input and one unresolved
    // The multiply initially sees one integer input and one unresolvedfcn8
    // recursive result, but the floating base case requires a floating mode.
    @Test
    public void testFcn5Float() {
        CodeGen code = new CodeGen("val _fact = { int x -> x <= 1 ? 1.0 : x*_fact(x-1); }; return _fact(arg);");
        code.parse().opto().typeCheck();
        assertEquals("1.0",  Eval2.eval(code, 0));
        assertEquals("1.0",  Eval2.eval(code, 1));
        assertEquals("2.0",  Eval2.eval(code, 2));
        assertEquals("6.0",  Eval2.eval(code, 3));
        assertEquals("24.0", Eval2.eval(code, 4));
    }

    @Test
    public void testFcn5Minus() {
        CodeGen code = new CodeGen("val _neg = { int x -> x ? -_neg(x-1) : 1; }; return _neg(arg);");
        code.parse().opto().typeCheck();
        assertEquals( "1", Eval2.eval(code, 0));
        assertEquals("-1", Eval2.eval(code, 1));
        assertEquals( "1", Eval2.eval(code, 2));
        assertEquals("-1", Eval2.eval(code, 3));
    }

    @Test
    public void testFcn6() {
        CodeGen code = new CodeGen(
"""
struct _S { int i; new _S = { int x -> i=x; }; };
val _newS = { int x -> return new _S(x); };
return _newS(1).i;
""");
        code.parse().opto().typeCheck();
        assertEquals("return 1;", code.print());
        assertEquals("1", Eval2.eval(code,  0));
    }

    // Only forward refs to direct function calls allowed
    @Test
    public void testFcn7() {
        CodeGen code = new CodeGen(
"""
val _f = {->1;};
val _g = {->2;};
if( arg ? _f : _g ) return 1;
return 2;
""");
        code.parse().opto().typeCheck();
        assertEquals("return 1;", code.print());
        assertEquals("1", Eval2.eval(code,  0));
    }

    @Test
    public void testFcn8() {
        CodeGen code = new CodeGen(
"""
{int -> int}? _i2i = null;
var _id = {{int->int} f-> return f;};
for(;;) {
    if (_i2i) return _i2i(arg);
    var _x = {int i-> return i;};
    arg = _x(3);
    _i2i = _id(_x);
}
""");
        code.driver(Phase.LocalSched);
        assertEquals("return Phi(Loop,arg,3);", code.print());
        assertEquals("3", Eval2.eval(code,  0));
    }

    @Test
    public void testFcn9() {
        CodeGen code = new CodeGen(
"""
{int -> int}? i2i = null;
for(;;) {
    if (i2i) return i2i(arg);
    var x = {int i-> return i;};
    arg = x(3);
}
""");
        code.driver(Phase.LocalSched);
        assertEquals("Stop[ return Top; return Parm_i(x,i64); ]", code.print());
        assertNull( Eval2.eval( code, 0 ) ); // Infinite loop, <clinit> never exits
    }


    @Test
    public void testFcn10() {
        CodeGen code = new CodeGen(
"""
struct _Person {
  int age;
};

val _fcn = { _Person?[] ps, int x ->
  val tmp = ps[x];
  if( ps[x] )
    ps[x].age++;
};

var ps = new _Person?[2];
ps[0] = new _Person;
ps[1] = new _Person;
return _fcn(ps,1);
""");
        code.driver(Phase.LocalSched);
        assertEquals("return .age;", code.print());
        assertEquals("0", Eval2.eval(code,  0));
    }

    // Function break
    @Test
    public void testErr1() {
        CodeGen code = new CodeGen(
"""
for(;;) {
    val f = { ->
        break;
    };
    f();
    return 2;
}
return 1;
""");
        try { code.parse().opto(); fail(); }
        catch( Exception e ) { assertEquals("No active loop for a break or continue",e.getMessage()); }
    }

    // Calling and inlining a null function
    @Test
    public void testErr2() {
        CodeGen code = new CodeGen(
"""
{int -> int}? !i2i = { int i -> return i; };
for(;;) {
    if (i2i(2) == arg) break;
    i2i = null;
}
return 0;
""");
        try { code.parse().opto().typeCheck(); fail(); }
        catch( Exception e ) { assertEquals("Might be null calling { ptr i64 -> i64 #[ 3]}?",e.getMessage()); }
    }


    @Test
    public void testErr3() {
        CodeGen code = new CodeGen(
"""
val f = { int i, int j -> return i+j; };
return f();
""");
        try { code.parse().opto().typeCheck(); fail(); }
        catch( Exception e ) { assertEquals("Expecting 3 arguments, but found 1",e.getMessage()); }
    }


    @Test
    public void testErr4() {
        CodeGen code = new CodeGen(
"""
struct S {
    {int} f = { -> x(); return 0; }; // Do not let fref x be a field
};
val x = { -> return 1; };
S? s = new S;
return s.x;
""");
        try { code.parse().opto().typeCheck(); fail(); }
        catch( Exception e ) { assertEquals("Accessing unknown field 'x' from '*Test.S'",e.getMessage()); }
    }

    @Test @Ignore
    public void testErr5() {
        CodeGen code = new CodeGen(
"""
val g = { ->
    g();
};
return g();
""");
        try { code.parse().opto().typeCheck(); fail(); }
        catch( Exception e ) { assertEquals("No defined return type",e.getMessage()); }
    }


    // Mutual recursion.  Fails without SCCP to lift the recursive return types.
    @Ignore @Test
    public void testFcnMutRec() {
        CodeGen code = new CodeGen(
"""
val is_even = { int x -> x ? is_odd (x-1) : true ; };
val is_odd  = { int x -> x ? is_even(x-1) : false; };
return is_even(arg);
""");
        code.parse().opto().typeCheck();
        assertEquals("Stop[ return #2; return Phi(Region,0,1,#2); ]", code.print());
        assertEquals("1", Eval2.eval(code, 0));
        assertEquals("0", Eval2.eval(code, 1));
        assertEquals("1", Eval2.eval(code, 2));
        assertEquals("0", Eval2.eval(code, 3));
    }

    // Forward ref to 'x' means that
    @Ignore @Test
    public void testForwardRef1() {
        CodeGen code = new CodeGen(
"""
struct S {
    { int } f = { -> return x(); };
};
val x = { -> return 1; };
S? s = null;
for(;;) {
    if (s) return s.f;
}
""");
        code.parse().opto().typeCheck().GCM();
        assertEquals("return 0;", code.print());
        assertEquals("0", Eval2.eval(code,  0));
    }

    // Inline hidden called more than once
    @Test
    public void testInline() {
        CodeGen code = new CodeGen(
"""
{int->int}?! _i2i_noInline = {int i->return i;};
{{int->int}->{int->int}}! _f2f = {{int->int} f->return f;};
val _o = _i2i_noInline;
if (arg) _i2i_noInline = null;
if (_i2i_noInline) return _i2i_noInline(arg);
return _f2f(_o)(1);
""");
        code.driver(Phase.LocalSched);
        assertEquals("Stop[ return Phi(Region,#2,#2); return Parm_i(_i2i_noInline,i64); ]", code.print());
        assertEquals("1", Eval2.eval(code,  2));
    }

    @Test
    public void testOperField() {
        CodeGen code = new CodeGen(
"""
struct _Person {
    int coffee_count;
};
_Person !p = new _Person;
p.coffee_count += 1;
return p.coffee_count;
""");
        code.driver(Phase.LocalSched);
        assertEquals("return 1;", code.print());
        assertEquals("1", Eval2.eval(code,  2));
    }

}