# Copyright 2018-2019, Carnegie Mellon University
# See LICENSE for details

update_type := function(vars, newType)
  vars.t := newType;
end;

Unique := (lst) ->
       FoldL(lst, (b, a)->Concat( b, When(not a in b, [a], [])), []);

SubstList:= function(expr, list_match, list_replace)
   local new_expr, index;
   new_expr := expr;
   for index in [1..Length(list_match)] do
     new_expr := SubstBottomUp(new_expr, @(2).cond(dc->dc=list_match[index]), x->list_replace[index]);
   od;
   return new_expr;
end;

Class(RulesStateHCOL, RuleSet, rec(inType := "SigmaSPL", outType := "SigmaSPL"));
RewriteRules(RulesStateHCOL, rec(
   Constant_PointWise := ARule(OLCompose, [@(1, PointWise, e->Length(Collect(e.op.expr, e.op.vars[1]))=0), @(2)], e->[@(1).val])
));

Class(RulesTerminateReductionHCOL, RuleSet, rec(inType := "SigmaSPL", outType := "SigmaSPL"));
RewriteRules(RulesTerminateReductionHCOL, rec(
     Reduction_GathH := ARule(OLCompose, [@(1, Reduction), @(2, GathH)],
       e->[ let(o:= @(1).val, i:= Ind(o.N),
        ISumReduction(i, o.N, o.op, o.idval, o.isSaturated, eT(@(2).val.N, add(i*@(2).val.stride, @(2).val.base)))) ]),

    Reduction_terminate := Rule(@(1, Reduction), e->let(o:= @(1).val, i := Ind(o.N),
        ISumReduction(i, o.N, o.op, o.idval, o.isSaturated, eT(o.N, i)))),

	Scat1Union_Gath1 := ARule(OLCompose, [@(1, ScatHUnion, e1->e1.n=1), @(2, GathH, e2->e2.n=1)],
		e->[ let(
		Print("Scat1Union_Gath1: ", @(1).val.N, " ", @(1).val.base, " ", @(2).val.N, " ", @(2).val.base,"\n"),
		OLCompose(
		  eUnion(@(1).val.N, @(1).val.base),
	      eT(@(2).val.N, @(2).val.base)
		)
	)]),

    ScatHUnion_GathH := ARule(OLCompose, [@(1, ScatHUnion), @(2, GathH)],
        e->[let( i:= Ind(@(1).val.n),
	  ISumUnion(i, @(1).val.n,
	    OLCompose(
	      eUnion(@(1).val.n, @(1).val.base+@(1).val.stride*i),
	      eT(@(2).val.n, @(2).val.base+@(2).val.stride*i)
	    )
	  )
	)]),

    GathH0_terminate := Rule(@(1, GathH, e->e.n=0), e->eT(0, -1)),
    GathH1_terminate := Rule(@(1, GathH, e->e.n=1), e->eT(@(1).val.N, @(1).val.base)),
    GathHN_terminate := Rule(@(1, GathH, e->e.n>1),
        e-> let(i := Ind(@(1).val.n),
	    	  ISumUnion(i, @(1).val.n,
		    OLCompose(
			eUnion(@(1).val.n, i),
			eT(@(1).val.N, @(1).val.base+@(1).val.stride*i))))),

    ScatHUnion0_terminate := Rule(@(1, ScatHUnion, e->e.n=0), e->eUnion(0,-1)),   #eUnion(@(1).val.N, @(1).val.base)),
    ScatHUnion1_terminate := Rule(@(1, ScatHUnion, e->e.n<=1), e->eUnion(@(1).val.N, @(1).val.base)),
    ScatHUnionN_terminate := Rule(@(1, ScatHUnion, e->e.n>1),
        e-> let(i := Ind(@(1).val.n),
	       ISumUnion(i, @(1).val.n,
 	         OLCompose(
			eUnion(@(1).val.N, @(1).val.base+@(1).val.stride*i),
			eT(@(1).val.n, i)
	         )
               ))
    ) ));

Class(RulesSumsHCOLv2a, RuleSet, rec(inType := "SigmaSPL", outType := "SigmaSPL"));
RewriteRules(RulesSumsHCOLv2a, rec(
    OLCompose_Assoc := ARule(OLCompose, [ @(1,OLCompose) ],  e -> @(1).val.children() ),
    OLCompose_PointWise_PointWise := ARule(OLCompose, [ @(1, PointWise), @(2, PointWise) ],
        e -> [ PointWise(@(1).val.N, Lambda(@(2).val.op.vars, SubstVars(Copy(@(1).val.op.expr),
         rec((@(1).val.op.vars[2].id) := @(2).val.op.vars[2], (@(1).val.op.vars[1].id) := @(2).val.op.expr)))) ]),
    GathH_GathH := ARule(OLCompose, [@(1, GathH), @(2, GathH)],
        e -> [GathH(@(2).val.N, @(1).val.n, @(1).val.base+@(2).val.base, @(1).val.stride*@(2).val.stride)]),
    Reduction_ScatHUnion := ARule(OLCompose, [ @(1, Reduction), @(2, ScatHUnion, e->e.n=1) ], e->[]),
    PointWise_BinOp := ARule(OLCompose, [@(1, PointWise, e->e.N=1), @(2, BinOp, e->e.N=1)],
        e -> [BinOp(1,Lambda(@(2).val.op.vars, @(1).val.op.at(@(2).val.op.expr, V(0))))]),
    PointWise_ISumUnion :=  ARule(OLCompose, [ @(1, PointWise), @(2, ISumUnion) ],
        e -> [ CopyFields(@(2).val, rec(
            _children :=  List(@(2).val._children, c -> OLCompose(@(1).val, c)),
            dimensions := [Rows(@(1).val), Cols(@(2).val)] )) ]),
    PointWise_ScatHUnion := ARule(OLCompose, [@(1,PointWise), @(2, ScatHUnion)],
        e -> let(i := Ind(@(2).val.dims()[2]), [@(2).val, PointWise(@(2).val.dims()[2],
            Lambda([@(1).val.op.vars[1], i], SubstVars(@(1).val.op.expr, rec((@(1).val.op.vars[2].id) := i*@(2).val.stride+@(2).val.base))))])),
    ISumXXX_YYY := ARule(OLCompose, [ @(1, [ISumUnion, ISumReduction]), @(2, [GathH, PointWise, Induction, eT]) ],
        e -> [ CopyFields(@(1).val, rec(
            _children :=  List(@(1).val._children, c -> OLCompose(c, @(2).val)),
            dimensions := [Rows(@(1).val), Cols(@(2).val)] )) ]),
    SUMUnion_GathH := ARule(OLCompose, [@(1, SUMUnion), @(2, GathH)],
       e ->[SUMUnion(List(@(1).val._children, c->OLCompose(c, @(2).val)))] ),
    ScatHUnion_SUMUnion := ARule(OLCompose, [ @(1, ScatHUnion), @(2, SUMUnion) ],
        e -> [SUMUnion(List(@(2).val._children, c->OLCompose(@(1).val,c)))] ),
    ScatHUnion_ISumUnion := ARule(OLCompose, [@(1, ScatHUnion), @(2, ISumUnion)],
        e -> [ISumUnion(@(2).val.var, @(2).val.domain, OLCompose(@(1).val, @(2).val._children[1]))]),
    ScatHUnion_ScatHUnion := ARule(OLCompose, [@(1, ScatHUnion), @(2, ScatHUnion)],
        e -> [ScatHUnion(@1.val.N, @(2).val.n, @(1).val.base+@(2).val.base, @(2).val.stride)] ),
    SUMUnion_Assoc := ARule(SUMUnion, [@(1, SUMUnion) ], e->@(1).val.children())
));

Class(RulesSumsHCOLv2b, RuleSet, rec(inType := "SigmaSPL", outType := "SigmaSPL"));
RewriteRules(RulesSumsHCOLv2b, rec(
    OLCompose_Assoc := ARule(OLCompose, [ @(1,OLCompose) ],  e -> @(1).val.children() ),
    OLCompose_PointWise_PointWise := ARule(OLCompose, [ @(1, PointWise), @(2, PointWise) ],
        e -> [ PointWise(@(1).val.N, Lambda(@(2).val.op.vars, SubstVars(Copy(@(1).val.op.expr),
         rec((@(1).val.op.vars[2].id) := @(2).val.op.vars[2], (@(1).val.op.vars[1].id) := @(2).val.op.expr)))) ]),
    GathH_GathH := ARule(OLCompose, [@(1, GathH), @(2, GathH)],
        e -> [GathH(@(2).val.N, @(1).val.n, @(1).val.base+@(2).val.base, @(1).val.stride*@(2).val.stride)]),
    ISumReduction_PointWise := ARule(OLCompose, [ @(1, ISumReduction), @(2, PointWise) ],
        e -> [ ISumReduction(@(1).val.var, @(1).val.domain, @(1).val.op, @(1).val.idval, @(1).val.isSaturated,
                OLCompose(@(1).val._children[1], @(2).val))]),

    eT_Pointwise := ARule(OLCompose, [@(1, eT, e->IsVar(e.base)), @(2, PointWise)],
         e -> let(i := Ind(1),
	     [PointWise(1, Lambda(
                SubstVars(Copy(@(2).val.op.vars), rec((@(2).val.op.vars[2].id):=i)),
                SubstVars(Copy(@(2).val.op.expr), rec((@(2).val.op.vars[2].id):=@(1).val.base)))), @(1).val])),
    ISumXXX_YYY := ARule(OLCompose, [ @(1, [ISumUnion, ISumReduction]), @(2, [GathH, PointWise, Induction, eT]) ],
        e -> [ CopyFields(@(1).val, rec(
            _children :=  List(@(1).val._children, c -> OLCompose(c, @(2).val)),
            dimensions := [Rows(@(1).val), Cols(@(2).val)] )) ]),
    eT_ISumUnion := ARule(OLCompose, [ @(1, eT), @(2, ISumUnion, e->ObjId(e._children[1]._children[1])=eUnion)],
        e -> [ Drop(SubstVars(Copy(@(2).val._children[1]._children), rec((@(2).val.var.id) := (@(1).val.base)) ), 1)]),
    eT_Induction := ARule(OLCompose, [@(1, eT), @(2, Induction)],
        e -> [Inductor(@(2).val.N, @(1).val.base, @(2).val.op, @(2).val.initval)]),
  PointWise_ISumUnion :=  ARule(OLCompose, [ @(1, PointWise), @(2, ISumUnion) ],
        e -> [ CopyFields(@(2).val, rec(
            _children :=  List(@(2).val._children, c -> OLCompose(@(1).val, c)),
            dimensions := [Rows(@(1).val), Cols(@(2).val)] )) ]),

   PointWise := ARule(OLCompose, [@(1, PointWise, e->Length(Collect(e.op.expr, e.op.vars[1]))=0), ...],
       e -> [@(1).val])
));

Class(RulesSumsHCOL, RuleSet, rec(inType := "SigmaSPL", outType := "SigmaSPL"));
RewriteRules(RulesSumsHCOL, rec(
    OLCompose_Assoc := ARule(OLCompose, [ @(1,OLCompose) ],  e -> @(1).val.children() ),
    OLCompose_PointWise_PointWise := ARule(OLCompose, [ @(1, PointWise), @(2, PointWise) ],
        e -> [ PointWise(@(1).val.N, Lambda(@(2).val.op.vars, SubstVars(Copy(@(1).val.op.expr),
         rec((@(1).val.op.vars[2].id) := @(2).val.op.vars[2], (@(1).val.op.vars[1].id) := @(2).val.op.expr)))) ]),
    PointWise_ISumUnion :=  ARule(OLCompose, [ @(1, PointWise), @(2, ISumUnion) ],
        e -> [ CopyFields(@(2).val, rec(
            _children :=  List(@(2).val._children, c -> OLCompose(@(1).val, c)),
            dimensions := [Rows(@(1).val), Cols(@(2).val)] )) ]),
    PointWise_ScatHUnion := ARule(OLCompose, [@(1,PointWise), @(2, ScatHUnion)],
        e -> let(i := Ind(@(2).val.dims()[2]), [@(2).val, PointWise(@(2).val.dims()[2],
            Lambda([@(1).val.op.vars[1], i], SubstVars(@(1).val.op.expr, rec((@(1).val.op.vars[2].id) := i*@(2).val.stride+@(2).val.base))))])),
    Reduction_ISumReduction :=  ARule(OLCompose, [ @(1, Reduction), @(2, ISumUnion) ],
        e -> [ ISumReduction(@(2).val.var, @(2).val.domain, @(1).val.op, @(1).val.idval, @(1).val.isSaturated,
               OLCompose(@(1).val, @(2).val._children[1]))]),
    Reduction_ScatHUnion := ARule(OLCompose, [ @(1, Reduction), @(2, ScatHUnion, e->e.n=1) ], e->[]),
    ISumReduction_PointWise := ARule(OLCompose, [ @(1, ISumReduction), @(2, PointWise) ],
        e -> [ ISumReduction(@(1).val.var, @(1).val.domain, @(1).val.op, @(1).val.idval, @(1).val.isSaturated,
                OLCompose(@(1).val._children[1], @(2).val))]),
    eT_Pointwise := ARule(OLCompose, [@(1, eT, e->IsVar(e.base)), @(2, PointWise)],
        e -> let(i := Ind(1),
            [PointWise(1, Lambda(
                SubstVars(Copy(@(2).val.op.vars), rec((@(2).val.op.vars[2].id):=i)),
                SubstVars(Copy(@(2).val.op.expr), rec((@(2).val.op.vars[2].id):=@(1).val.base)))), @(1).val])),
    ISumXXX_YYY := ARule(OLCompose, [ @(1, [ISumUnion, ISumReduction]), @(2, [GathH, PointWise, Induction, eT]) ],
        e -> [ CopyFields(@(1).val, rec(
            _children :=  List(@(1).val._children, c -> OLCompose(c, @(2).val)),
            dimensions := [Rows(@(1).val), Cols(@(2).val)] )) ]),
    GathH_GathH := ARule(OLCompose, [@(1, GathH), @(2, GathH)],
        e -> [GathH(@(2).val.N, @(1).val.n, @(1).val.base+@(2).val.base, @(1).val.stride*@(2).val.stride)]),
    PointWise_BinOp := ARule(OLCompose, [@(1, PointWise, e->e.N=1), @(2, BinOp, e->e.N=1)],
        e -> [BinOp(1,Lambda(@(2).val.op.vars, @(1).val.op.at(@(2).val.op.expr, V(0))))]),
    eT_Induction := ARule(OLCompose, [@(1, eT), @(2, Induction)],
        e -> [Inductor(@(2).val.N, @(1).val.base, @(2).val.op, @(2).val.initval)]),
    ScatHUnion_ScatHUnion := ARule(OLCompose, [@(1, ScatHUnion), @(2, ScatHUnion)],
        e -> [ScatHUnion(@1.val.N, @(2).val.n, @(1).val.base+@(2).val.base, @(2).val.stride)] )
));

Class(ToVectors, HierarchicalVisitor, rec(
    __call__ := meth(arg)
    local res;
        res := ApplyFunc(arg[1].visit, arg{[2..Length(arg)]});
       return res;
    end,
    PointWise := (self, o, opts) >> let(
      new_i := Ind(4), new_var := var.fresh_t("r", TReal), times := (o.N - Mod(o.N, 4))/4,
      DirectSum(
        List([0..times-1], i->PointWise(4, Lambda([new_var, new_i], new_var*i)) ),
	PointWise(Mod(o.N, 4), o.op)
      )
    )
));

max_value := var("max_val", TReal);
max_err := var("max_err", TReal);
Class(RulesErrorHCOL, RuleSet, rec(inType := "iCode", outType := "iCode"));
RewriteRules(RulesErrorHCOL, rec(
  assign_const := Rule([@(1, assign), @(2, var), @(3, Value, v->v<>0)],
					e->assign(@(2).val, @(3).val+@(3).val*max_err) ),
  assign_nth := Rule([@(1, assign), @(2, var), @(3, nth)],
					e->assign(@(2).val, max_value+max_value*max_err) )
));

Class(RulesCodeHCOL, RuleSet, rec(inType := "iCode", outType := "iCode"));
RewriteRules(RulesCodeHCOL, rec(
#only assignments to state but state is not used

	simplify_state := Rule(@@(1, decl, (e, cx)->(
	  let(
	  Length(Collect(e.cmd, decl))=0 and
	  IsBound(cx.opts.state) and
	  Length(Collect(e.cmd, @(1,nth).cond(j->IsVar(j.loc) and j.loc.id=cx.opts.state.id))) > 0 and
	  Length(Collect(e.cmd, @(1,nth).cond(j->IsVar(j.loc) and j.loc.id=cx.opts.state.id)))=Length(Collect(e.cmd, @(1, assign).cond(i->ObjId(i.loc)=nth and i.loc.loc.id = cx.opts.state.id))) ))),
	  (e, cx)-> let(Print("Remove State:\n", @@(1).val, "\n"), SubstTopDown(@@(1).val, [@(2, assign), @(3, nth).cond(i->i.loc.id=cx.opts.state.id), @(4)], j->skip()))),

    chain_chain := Rule(@(1, chain, e->ObjId(e.cmds[1])=chain),
        e -> let(chain(Concat(@(1).val.cmds[1].cmds, Drop(@(1).val.cmds, 1))))),

	#replaces TArray of length one with a scalar variable
    scalarize1 := Rule(@(1, decl, e->Length(Filtered(e.vars, i->ObjId(i.t)=TArray and i.t.size=1))>=1),
        e->let(
			oldvar := Filtered(e.vars, i->ObjId(i.t)=TArray and i.t.size=1)[1],
			Print("Scalarizing ", oldvar, "\n"),

			tmp := Collect(@(1).val, [@(5, assign), @(6).cond(ee->IsNth(ee) and ee.loc.id=oldvar.id), @(7)]),
			newType := UnifyTypes(List(tmp, i->i.exp.t)),
			newvar := var.fresh_t("s", newType),
			Print("Scalarizing ", oldvar, " with ", newvar, " of type: ", newType, "\n"),

            decl(
                SubstVars(@(1).val.vars, rec((oldvar.id):=newvar)),
                SubstTopDown(@(1).val.cmd, [@(2, nth), @(3).cond(e->IsVar(e) and e.id=oldvar.id), @(4).cond(e->IsValue(e) and e.v = 0)],
                    i->newvar)
            ))),

	#replace each fixed index array reference (i.e. nth(X, i), where i is a constant value) with a variable
	#This can only be done if there are no array references to X that is variable (this is to prevent potential aliasing)
    scalarize_const := Rule(@(1, decl, e->Length(Filtered(Filtered(e.vars, j->IsArray(j.t)), k->ForAll(Collect(e.cmd, [@(2, nth), @(3, var, u->u.id=k.id), @(4)]), j->IsValue(j.idx))))>=1),
        e->let(
			oldvar := Filtered(Filtered(e.vars, j->IsArray(j.t)),   #find all array variables
	           k->ForAll(Collect(e.cmd, [@(2, nth), @(3, var, u->u.id=k.id), @(4)]), j->IsValue(j.idx)))[1],   #make sure only const index

            newvars := List([1..oldvar.t.size], j->var.fresh_t("q", oldvar.t.t)),
            cmmd := @(1).val.cmd,
            lst := List([0..Length(newvars)-1], i->[i, newvars[i+1]]),
            f := (c,l)->SubstTopDown(c, [@(2, nth), @(3).cond(k->IsVar(k) and k.id=oldvar.id), @(4).cond(j->IsValue(j) and j.v = l[1])], i->l[2]),
            Print("Scalarize_const: ",oldvar,":: ", newvars, "\n"),
            decl(Concat(Filtered(@(1).val.vars, j->j.id <> oldvar.id), newvars), FoldL(lst, f, cmmd))
        )),

    chain_xyz_chain := ARule(chain, [@(1), @(2, chain)],
        e->[chain(Concat([@(1).val], @(2).val.cmds))]),
    chain_chain_xyz := ARule(chain, [@(1, chain), @(2)],
        e->[chain(Concat(@(1).val.cmds), [@(2).val])]),

    chain_creturn := ARule(chain, [@(1, creturn), @(2)],
        e->[chain(Concat([@(2).val], [@(1).val]))] ),

    loop_pull_cond_1 := Rule([@(1, loop), @(2), @(3), [@(4, chain),  [@(5, assign), @(6,var),
            [@(7,cond), @(8,eq, e->e.args[1]=@(2).val and e.args[2].v=0), @(9), @(10)]],...]],
        e->let(
	   chain(assign(@(6).val, @(9).val), loop(@(1).val.var, @(1).val.range, chain([assign(@(6).val, @(10).val)]::Drop(@(4).val.cmds,1)))))),

	   copyprop_var_val := Rule(@(1, decl, e->let(
		   #only propagate values if the variable is SSA and only if it is within a basic block
		   cnds := Collect(e.cmd, [assign, @(2, var, ee->ee in e.vars), @(3, Value)]),
		   actual_cnds := Filtered(cnds, i->Length(Collect(e.cmd, [assign, @(8, var, ee->ee.id = i.loc.id), @(9)]))=1),
		   #Print("CopyProp: ", cnds, "\n\n", e.cmd, "\n\n", actual_cnds, "\n"),
		   Length(actual_cnds) >= 1 and Length(Collect(e.cmd, [assign, @(4, var, ee->ee.id = actual_cnds[1].loc.id), @(5)]))=1)),
		e->let(
			all_assigned_value := Collect(@(1).val.cmd, [assign, @(6, var, ee->ee in e.vars), @(7, Value)]),
			assigned_var_value := Filtered(all_assigned_value, i->Length(Collect(@(1).val.cmd, [assign, @(8, var, ee->ee.id = i.loc.id), @(9)]))=1)[1],
		   #Print("CopyProp_value: ", assigned_var_value, "\n"),
			decl(Filtered(@(1).val.vars, j->j.id<>assigned_var_value.loc.id),
				 SubstVars(@(1).val.cmd, rec((assigned_var_value.loc.id) := assigned_var_value.exp)))
		)
       ),

	   copyprop_var_var := Rule(@(1, decl, e->let(
		   #only propagate var if the assigned variable is SSA
           cnds := Collect(e.cmd, [assign, @(2, var, ee->ee in e.vars), @(3, var, e->e.id<>@(2).val.id)]),
           Length(cnds)>=1 and cnds[1].loc in e.vars and Length(Collect(e.cmd, [assign, @(4, var, ee->ee.id=cnds[1].loc.id), @(5)]))=1)),
        e->let(
			   acnds := Collect(@(1).val.cmd, [assign, @(7, var, ee->ee in e.vars), @(8, var, e->e.id<>@(7).val.id)]),
			   atests := Collect(e.cmd, [assign, @(9, var, ee->ee.id=acnds[1].loc.id), @(5)]),
			   a := atests[1],
			   decl(Filtered(@(1).val.vars, j->j <>a.loc), SubstVars(@(1).val.cmd, rec((a.loc.id) := a.exp))))),

    drop_selfassign_var := ARule(chain, [@(1, assign, e->IsVar(e.loc) and IsVar(e.exp) and e.exp=e.loc), @(2)],
        e->let(
		[@(2).val])),

    drop_selfassign2 := ARule(chain, [@(1, assign, e->e.exp=e.loc), @(2)],
        e->[@(2).val]),

	#this drops unused variables i.e. variables that are declared but not used.
    drop_unused_var := Rule(@(1, decl, e-> let(v := Collect(e.cmd, var), Filtered(e.vars, k->not k in v))<> []),
        e->let(v := Collect(@(1).val.cmd, var),
            decl(Filtered(@(1).val.vars, k->k in v), @(1).val.cmd))),

	drop_unused_assign := Rule(@(1, decl, e->
			Filtered(List(e.vars, v->Collect(e.cmd, v)), l->Length(l)=1)<>[]) ,
	e->let(
      unused := Filtered(List(e.vars, v->Collect(e.cmd, v)), l->Length(l)=1)[1],
	  empty_list := List([1..Length(unused)], i->skip()),
	  target_list := Filtered(Collect(@(1).val.cmd, assign), a->a.loc in unused),
	  Print("Unused assignements:\n", target_list, "\n\n"),
	  decl(Filtered(@(1).val.vars, k->not k in unused), SubstList(@(1).val.cmd, target_list, empty_list))
	)),

    loop_decl := Rule(@(1, loop, e->ObjId(e.cmd)=decl and Length(Collect(e.cmd, decl))=1 and Length(Collect(e.cmd, chain))<=1),
        e -> decl(@(1).val.cmd.vars, loop(@(1).val.var, @(1).val.range, @(1).val.cmd.cmd))),

    chain_xyz_decl := ARule(chain, [@(1), @(2, decl)],
        e->[decl(@(2).val.vars, chain(@(1).val, @(2).val.cmd))]),

    chain_decl := Rule(@(1, chain, e->ObjId(e.cmds[1])=decl),
        e -> decl(@(1).val.cmds[1].vars, chain(Concat([@(1).val.cmds[1].cmd], Drop(@(1).val.cmds, 1))))),
    decl_decl := Rule(@(1, decl, e->ObjId(e.cmd)=decl),
        e -> decl(Set(Concat(@(1).val.vars, @(1).val.cmd.vars)), @(1).val.cmd.cmd)),
));

Class(RulesHCOLnoAbsMax, RuleSet, rec(inType := "iCode", outType := "iCode"));
RewriteRules(RulesHCOLnoAbsMax, rec(
    abs_cond := Rule([@(1, assign), @(2), @(3, abs)],
        e->let(s := var.fresh_t("w", @(3).val.t), decl(s, chain(assign(s, @(3).val.args[1]), assign(@(1).val.loc, cond(geq(s, V(0)), s, neg(s))))))),
    max_cond := Rule(@(1, max),
        e->cond(geq(@(1).val.args[1], @(1).val.args[2]), @(1).val.args[1], @(1).val.args[2])),
));

Class(RulesUnrollHCOL, RuleSet, RulesStrengthReduce, RulesCodeHCOL, rec(inType := "iCode", outType := "iCode"));
RewriteRules(RulesUnrollHCOL, rec(
    unroll_wo_decl := Rule(@@(1, [loopn, loop],
       (e, cx) -> Length(Collect(e.cmd, loop))=0 and
                  Length(Collect(e.cmd, decl))=0 and
	          e.var.range <= When(IsBound(cx.opts) and IsBound(cx.opts.globalUnrolling),
		       	                                           cx.opts.globalUnrolling,
                                              5)),
       (e,cx)->let(i := @@(1).val.var,
	               rng := @@(1).val.range,
		       vvars := Filtered(Collect(@@(1).val.cmd, var), j->j.t in [TReal, TDouble]),

           chain(
	      List(rng, j->let(
		    ssvars := FoldL([[i, V(j)]], (b, a)->CopyFields(rec((a[1].id):= a[2]), b), rec()),
		    new_cmd := SubstVars(Copy(@@(1).val.cmd), ssvars),
 		    new_cmd ))
    ))) ,
    unroll_w_decl := Rule(@@(1, [loop, loopn],
           (e,cx)->Length(Collect(e.cmd, loop))=0 and
				   Length(Collect(e.cmd, decl))=1 and
	           e.var.range <= When(IsBound(cx.opts) and IsBound(cx.opts.globalUnrolling),
           	                                          cx.opts.globalUnrolling,
                                              5)),
           (e,cx)->let(i := @@(1).val.var,
	               rng := @@(1).val.range,
		       vvars := Filtered(Collect(@@(1).val.cmd, var), j->j.t in [TReal, TDouble]),

		       localvar := Collect(@@(1).val.cmd, decl)[1].vars,

           chain(
	      List(rng, j->let(

	      		        nvars :=FoldL(localvar, (b,a)->Concat([[a, var.fresh_t("u", a.t)]], b), [[i, V(j)]] ),
	      		        ssvars := FoldL(nvars, (b, a)->CopyFields(rec((a[1].id):= a[2]) , b), rec()),
						svars := FoldL(nvars, (b, a)->Concat( When(IsVar(a[2]) or a[2] in b, [a[2]], []), b), []),

				new_cmd := SubstVars(Copy(@@(1).val.cmd.cmd), ssvars),

				decl(svars, new_cmd) )

                                #decl(svars, new_cmd ) )
             ))
	  )
   )
));

RewriteRules(RulesStrengthReduce, rec(
	addsub00 := Rule([@(1, addsub_2x64f), @(2).cond(e->Cond(IsValue(e), isValueZero(e), e=0)), _0],
		e->@(2).val)
));

Class(RulesUnifyType, RuleSet, rec(inType := "iCode", outType := "iCode"));
RewriteRules(RulesUnifyType, rec(
    unify_value := Rule(@(1,assign, e-> IsValue(e.exp) and IsVar(e.loc) and e.loc.t <> e.exp.t),
                        e->let(
                            newtype := UnifyTypes([@(1).val.exp.t, e.loc.t]),
                            Print("Unifying Type from a value to ", newtype,"\n"),
                            #if the newtype is the same as the loc, then need to cast the value
                            #else error, upcast the type of the variable since we are losing info.
                            Cond(newtype=e.loc.t,
                            assign(@(1).val.loc, tcast(newtype, @(1).val.exp)),
                            let(
                            update_type(@(1).val.loc, newtype),
                            assign(@(1).val.loc, @(1).val.exp)
                            ))
                        )),
	unify_vartype := Rule(@(1,assign,e->IsVar(e.loc) and not IsValue(e.exp) and e.loc.t <> e.exp.computeType()),
	   e->let(
		  oldtype := @(1).val.exp.t,
		  newtype := @(1).val.exp.computeType(),
		  Print("Unifying Type from ", oldtype, " to ", newtype,"\n"),
		  update_type(@(1).val.loc, newtype),
		  assign(@(1).val.loc, @(1).val.exp)
	   )),
));
RulesTypeHCOL := CopyFields(MergedRuleSet(RulesUnifyType),
    rec(inType:="iCode", outType := "iCode"));

RulesOrigCodeUnrollHCOL := CopyFields(MergedRuleSet(RulesUnrollHCOL, RulesStrengthReduce, RulesCodeHCOL, RulesHCOLnoAbsMax),
    rec(inType:="iCode", outType := "iCode"));

RulesCodeUnrollHCOL := CopyFields(MergedRuleSet( RulesUnrollHCOL, RulesStrengthReduce,  RulesCodeHCOL, RulesHCOLnoAbsMax),
    rec(inType:="iCode", outType := "iCode"));

RulesCodeNoUnrollHCOL := CopyFields(MergedRuleSet(RulesStrengthReduce, RulesCodeHCOL, RulesHCOLnoAbsMax),
    rec(inType:="iCode", outType := "iCode"));

#changes behaviour so that Spiral outputs abs(a) and max(a,b) C-code instead of "(a >= b) ? a : b"
RulesCodeUnrollHCOLuseAbsMaxSR := CopyFields(MergedRuleSet(RulesUnrollHCOL, RulesStrengthReduce, RulesCodeHCOL),
    rec(inType:="iCode", outType := "iCode"));


# Copyright (c) 2018-2021, Carnegie Mellon University
# See LICENSE for details

# Higher-order functions support
#

# Functional expression base class
Class(FuncExp, Exp, rec(
    isFuncExp := true,

    at := arg >> let(
	self := arg[1],
	Cond(Length(arg)=2 and IsList(arg[2]),
	     ApplyFunc(fcall, [self] :: arg[2]),
	     ApplyFunc(fcall, arg))),

    mkVars := self >>
        List(DropLast(self.computeType().params, 1), x->var.fresh_t("q", x)),

    getRankType := (self, r) >> let(
	t := self.computeType(),
	rank := Length(t.params)-2,
	Cond(r > rank,
	         TDummy,
	     t.params[ Length(t.params) - 1 - r ])
    )
));

Function.mkVars      := FuncExp.mkVars;
Function.getRankType := FuncExp.getRankType;

Symbolic.mkVars      := FuncExp.mkVars;
Symbolic.getRankType := FuncExp.getRankType;

IsFuncExp := o -> IsRec(o) and
    ((IsBound(o.isFuncExp) and o.isFuncExp) or (IsBound(o.t) and ObjId(o.t)=TFunc));

# Same as IsFuncExp(<o>) but returns false for param's with t = TFunc
# (these are considered 'trivial')
IsNonTrivFuncExp := o -> IsRec(o) and (IsBound(o.isFuncExp) and o.isFuncExp);

# !!! Below is no longer needed/used, autolib handles this via Parametrizer
# These are hacks for dealing with types which are integers like 2
# these really mean the range of an integer variable
# _convType converts 2 to TInt not to confuse other parts of the system
#

# !!! Below is no longer needed/used, autolib handles this via Parametrizer
#_convType := t -> Cond(IsSymbolic(t) or IsInt(t) or IsValue(t), TInt,
#    SubstBottomUp(Copy(t), TFunc,
#        e -> ApplyFunc(TFunc, List(e.params, p -> When(IsInt(p) or IsSymbolic(p) or IsValue(p), TInt, p)))));
_convType := t->t;

# !!! Below is no longer needed/used, autolib handles this via Parametrizer
#_convTypeExp := t ->
#    SubstBottomUp(Copy(t), TFunc,
#        e -> ApplyFunc(TFunc, List(e.params, p -> When(IsInt(p) or IsSymbolic(p) or IsValue(p), TInt, p))));
_convTypeExp := t->t;

#F fcurry(<func>, <pos>, <arg>) - symbolic representation of function currying
#F     Given a function with n arguments returns a functino with n-1 arguments
#F     where <arg> is plugged into position <pos>
#F
#F Example: p := param(TFunc(TInt, 4, TComplex), "p");  # ie. rank-1 4-elt diagonal func
#F          i1 := Ind(4);
#F          fcurry(p, 1, i1).t;
#F            => TFunc(4, TComplex);
#F
Class(fcurry, FuncExp, rec(
    computeType := self >> let(
        n := self.args[2],
        ct := self.args[1].computeType(),
        Checked(IsValue(n), ObjId(ct)=TFunc, n.v < Length(ct.params),
            ApplyFunc(TFunc, ListWithout(ct.params, n.v)))),

    eval := self >> let(
        torig := self.args[1].t.params,
        tcurried := self.t.params,
        n := self.args[2].v,
        vars := List([1..Length(tcurried)-1], i -> var.fresh_t("q", tcurried[i])),
        plugin := List([1..Length(torig)-1], i -> Cond(i<n, vars[i], i=n, self.args[3],vars[i-1])),
        Lambda(vars, ApplyFunc(fcall, Concatenation([self.args[1].eval()], plugin)))),

    at := (self, i) >> self.eval().at(i)
));

#F flift(<func>, <t>) - symbolic representation of function "lifting"
#F
#F  Given a function with n arguments returns a function with n+1 arguments,
#F  extra argument is of type <t> and is the "one but last" argument of new function.
#F  It is the ignored when computing the value of the function.
#F
#F  The purpose of flift is to increase the rank of functions. It implicitly creates
#F  a new inner loop, which loop variable is ignored. This is needed for rewrite rules
#F  like GT(T, ..) * Gath(f) -> GT(T*Gath(f), ...), where f is pulled into the loop.
#F
Class(flift, FuncExp, rec(
    computeType := self >> let(ct := self.args[1].computeType(), Checked(
        ObjId(ct)=TFunc, IsType(self.args[2]),
        let(tt := ct.params, n := Length(tt),
        ApplyFunc(TFunc, Concatenation(tt{[1..n-2]}, [self.args[2]], tt{[n-1..n]}))))),

    # NOTE: define what exactly .lambda() does in these cases
    lambda := self >> self.eval(),

    eval := self >> let(vars := List(DropLast(self.t.params,1), x->var.fresh_t("q", x)), n := Length(vars),
        Lambda(vars, ApplyFunc(fcall, Concatenation([self.args[1].eval()], ListWithout(vars, n-1)))))
));

#F fsplit(<func>, <loopid>, <inner_its>, <outer_its>)
Class(fsplit, FuncExp, rec(
    computeType := self >> let(ct := self.args[1].computeType(), Checked(
        ObjId(ct)=TFunc, IsValue(self.args[2]), self.args[2].t=TInt,
        let(tt := ct.params, loopid := self.args[2].v, pos := Length(tt)-1-loopid,
            ApplyFunc(TFunc, Concatenation(tt{[1..pos-1]}, [tt[pos], tt[pos]], tt{[pos+1..Length(tt)]}))))),

    # NOTE: define what exactly .lambda() does in these cases
    lambda := self >> self.eval(),

    eval := self >> let(
        vars      := List(DropLast(self.computeType().params, 1), x->var.fresh_t("q", x)),
        loopid    := self.args[2].v,
        inner_its := self.args[3],
        pos       := Length(vars)-loopid,
        callargs  := vars{[1..pos-2]} ::
	             Cond(vars[pos].t=TDummy, [0], [inner_its * vars[pos-1] + vars[pos]]) ::
                     vars{[pos+1..Length(vars)]},
        Lambda(vars, ApplyFunc(fcall, Concatenation([self.args[1].eval()], callargs))))
));

#F frotate(<func>, <n>)
#F   switch the order of loops (=ranks), by making <n>-th loop innermost
Class(frotate, FuncExp, rec(
    computeType := self >> let(ct := self.args[1].computeType(), nn := self.args[2], Checked(
        ObjId(ct)=TFunc, IsValue(nn), nn.t=TInt,
	let(tt := ct.params, rank := Length(tt)-2, n := nn.v, pos := rank+1-n,
            Cond(rank=0 or (rank=1 and n=1), ct,
		 n > rank,
		            ApplyFunc(TFunc, tt{[1..rank]} :: [TDummy] :: [tt[rank+1], tt[rank+2]]),
		 # else
                            ApplyFunc(TFunc, tt{[1..pos-1]} :: tt{[pos+1..rank]} :: [tt[pos], tt[rank+1], tt[rank+2]]))))),

    # NOTE: define what exactly .lambda() does in these cases
    lambda := self >> self.eval(),

    eval := self >> let(f := self.args[1], rank := f.rank(), n := self.args[2].v, pos := rank+1-n,
	Cond(rank=0 or (rank=1 and n=1), f,
	     n > rank,
             let(vars := List(DropLast(self.computeType().params, 1), x->var.fresh_t("q", x)),
		 Lambda(vars, ApplyFunc(fcall, [f] :: vars{[Length(vars)-1-rank..Length(vars)-2]} :: [Last(vars)]))),
	     # else
             let(vars := List(DropLast(self.computeType().params, 1), x->var.fresh_t("q", x)),
		 Lambda(vars, ApplyFunc(fcall, [f] :: vars{[1..pos-1]} :: [vars[rank]] :: vars{[pos..rank-1]} :: [vars[rank+1]])))
	))
));

## Ranked functions support
## Ranked functions == functions with implicit dependencies on loop variables
##
## These are used in paradigms.common.GT and autolib.*
##
## NOTE: get rid of _, these functions are not private, but public exports
##
_rankManip := (obj, newfunc) >>
    SubstTopDownNR(Copy(obj), @.cond(e->IsFunction(e) or IsFuncExp(e)), e -> newfunc(e));

_rank     := o -> Cond(
    IsList(o), Maximum0(List(o, _rank)),
    not IsRec(o) or not IsBound(o.rank), 0, o.rank());
_upRank   := o -> Cond(
    IsList(o), List(o, _upRank),
    not IsRec(o) or not IsBound(o.upRank), o, o.upRank());
_upRankNeq0   := o -> Cond(
    IsList(o), List(o, _upRankNeq0),
    not IsRec(o) or not IsBound(o.upRank) or o.rank()=0, o, o.upRank());
_upRankBy := (o,n) -> Cond(IsList(o), # NOTE: =0 ??
    List(o, x->_upRankBy(x,n)),
    not IsRec(o) or not IsBound(o.upRankBy) or o.rank()=0, o, o.upRankBy(n));
_downRank := (o,loopid,ind) -> Cond(
    IsList(o), List(o, x->_downRank(x, loopid, ind)),
    not IsRec(o) or not IsBound(o.downRank), o, o.downRank(loopid, ind));
_downRankFull := (o,inds) -> Cond(
    IsList(o), List(o, x->_downRankFull(x, inds)),
    not IsRec(o) or not IsBound(o.downRankFull), o, o.downRankFull(inds));
_split := (o, loopid, iits, oits) -> Cond(
    IsList(o), List(o, x->_split(x, loopid, iits, oits)),
    not IsRec(o) or not IsBound(o.split), o, o.split(loopid, iits, oits));
_rotate := (o, n) -> Cond(
    IsList(o), List(o, x->_rotate(x, n)),
    not IsRec(o) or not IsBound(o.rotate), o, o.rotate(n));

_rch_rank      := self >> Maximum0(List(self.rChildren(), _rank));
_rch_upRank    := self >> self.from_rChildren(List(self.rChildren(), _upRank));
_rch_upRankBy  := (self, n) >>  self.from_rChildren(List(self.rChildren(), c->_upRankBy(c,n)));
_rch_downRank  := (self, loopid, ind) >>
    self.from_rChildren(List(self.rChildren(), c->_downRank(c,loopid,ind)));
_rch_downRankFull := (self, inds) >>
    self.from_rChildren(List(self.rChildren(), c->_downRankFull(c,inds)));
_rch_split     := (self, loopid, iits, oits) >>
    self.from_rChildren(List(self.rChildren(), c->_split(c,loopid, iits, oits)));
_rch_rotate    := (self, n) >>
    self.from_rChildren(List(self.rChildren(), c->_rotate(c, n)));

Symbolic.domain := self >> Checked(ObjId(self.t)=TFunc, Length(self.t.params) > 1,
    self.computeType().params[Length(self.t.params)-1]);

Symbolic.range := self >> Checked(ObjId(self.t)=TFunc, Length(self.t.params) > 1,
    self.computeType().params[Length(self.t.params)]);

Symbolic.rank := self >> Cond(ObjId(self.t)=TFunc and Length(self.t.params) > 1,
    Length(self.t.params) - 2,
    _rch_rank(self));

Symbolic.at := (self, vars) >> Checked(ObjId(self.t)=TFunc, self.lambda().at(vars));

Symbolic.lambda := self >> Checked(ObjId(self.t)=TFunc, let(
    selft := self.computeType(),
    vars := List([1..self.rank()+1], x -> let(t:=selft.params[x],
            Cond(IsType(t), var.fresh_t("w", t), var.fresh("w", TInt, t)))),
    Lambda(vars, ApplyFunc(fcall, Concatenation([self], vars)))));

Symbolic.upRank := self >> Cond(ObjId(self.t)<>TFunc, _rch_upRank(self), flift(self, TDummy));

Symbolic.upRankBy := (self, n) >> Cond(
    ObjId(self.t)<>TFunc, _rch_upRankBy(self,n),
    Checked(IsPosInt0(n), FoldL([1..n], (f, i) -> f.upRank(), self)));

Symbolic.downRankFull := (self, inds) >> Cond(
    ObjId(self.t)<>TFunc, _rch_downRankFull(self, inds),
    FoldL(Reversed([1..Minimum(self.rank(), Length(inds))]), (f, i) -> f.downRank(i, inds[i]), self));

Symbolic.downRank := (self, loopid, ind) >> let(rank := self.rank(), Cond(
    loopid > rank, self,
    ObjId(self.t)<>TFunc, _rch_downRank(self, loopid, ind),
    fcurry(self, rank+1-loopid, ind)));  # !! loop variables are ordered with decreasing rank, highest-rank (outermost) is first var

Symbolic.split := (self, loopid, inner_its, outer_its) >> Cond(
    loopid > self.rank(), self,
    ObjId(self.t)<>TFunc, _rch_split(self, loopid, inner_its, outer_its),
    fsplit(self, loopid, inner_its, outer_its));

Symbolic.rotate := (self, n) >> Cond(
    self.rank() <= 1 and n <= 1, self,
    ObjId(self.t)<>TFunc, _rch_rotate(self, n),
    frotate(self, n));

# NOTE: below is really a hack, since the first argument might also have ranked
# things, although this is not obvious at first glance, an example of such first
# argument would be fcurry(func, lambdaWrap(rank-n function))

fcall.rank := self >> Cond(ObjId(self.t)=TFunc and Length(self.t.params) > 1,
    Length(self.t.params) - 2,
    _rank(Drop(self.args, 1)));

fcall.upRank := self >> Cond(ObjId(self.t)<>TFunc,
    ApplyFunc(fcall, [self.args[1]] :: _upRank(Drop(self.args, 1))),
    flift(self, TDummy));

fcall.upRankBy := (self, n) >> Cond(ObjId(self.t)<>TFunc,
    ApplyFunc(fcall, [self.args[1]] :: _upRankBy(Drop(self.args, 1), n)),
    Checked(IsPosInt0(n), FoldL([1..n], (f, i) -> f.upRank(), self)));

fcall.downRankFull := (self, inds) >> Cond(ObjId(self.t)<>TFunc,
    ApplyFunc(fcall, [self.args[1]] :: _downRankFull(Drop(self.args, 1), inds)),
    FoldL(Reversed([1..Minimum(self.rank(), Length(inds))]), (f, i) -> f.downRank(i, inds[i]), self));

fcall.downRank := (self, loopid, ind) >> let(rank := self.rank(), Cond(
    loopid > rank, self,
    ObjId(self.t) <> TFunc,
        ApplyFunc(fcall, [self.args[1]] :: _downRank(Drop(self.args, 1), loopid, ind)),
    # else
    fcurry(self, rank+1-loopid, ind)));  # !! loop variables are ordered with decreasing rank, highest-rank (outermost) is first var

fcall.split := (self, loopid, inner_its, outer_its) >> Cond(
    loopid > self.rank(), self,
    ObjId(self.t) <> TFunc,
        ApplyFunc(fcall, [self.args[1]] :: _split(Drop(self.args, 1), loopid, inner_its, outer_its)),
    fsplit(self, loopid, inner_its, outer_its));

fcall.rotate := (self, n) >> Cond(
    self.rank() <= 1 and n <= 1, self,
    ObjId(self.t) <> TFunc,
        ApplyFunc(fcall, [self.args[1]] :: _rotate(Drop(self.args, 1), n)),
    frotate(self, n));

Function.rank      := _rch_rank;
Function.upRank    := _rch_upRank;
Function.upRankBy  := _rch_upRankBy;
Function.downRank  := _rch_downRank;
Function.downRankFull := _rch_downRankFull;
Function.split     := _rch_split;
Function.rotate    := _rch_rotate;
Function.computeType := self >> self.lambda().t;

Lambda.rank      := Symbolic.rank;
Lambda.upRank    := Symbolic.upRank;
Lambda.upRankBy  := Symbolic.upRankBy;
#Lambda.downRank  := Symbolic.downRank;  downRank now defined in lambda.gi
Lambda.downRankFull := Symbolic.downRankFull;
Lambda.split     := Symbolic.split;
Lambda.rotate    := Symbolic.rotate;

# ind(<range>, <n>) - "nameless" reference to a loop index of n-th inner most loop
#            (eg. ind(1) is inner most, ind(n) is outermost in n-loop nest)
#            loop counter runs from 0..range-1
Class(ind, Loc, rec(
    __call__ := (self, range, n) >> Checked(IsInt(n),
	WithBases(self, rec(operations := ExpOps, range:=range, n:=n))),
    print := self >> Print(self.name, "(", self.range, ", ", self.n, ")"),
    rChildren := self >> [self.range, self.n],
    rSetChild := rSetChildFields("range", "n"),
    t := TInt,
    eval := self >> self,
    can_fold := False,

    upRank := self >> ObjId(self)(self.range, self.n+1),

    split := (self, loopid, inner_its, outer_its) >>
        Cond( loopid > self.n, self,
              loopid < self.n, ObjId(self)(self.range, self.n+1),
              ObjId(self)(inner_its, self.n) + inner_its*ObjId(self)(outer_its, self.n+1)),
    rotate := (self, n) >>
        Cond( n < self.n, self,
              n > self.n, ObjId(self)(self.range, self.n+1),
              ObjId(self)(self.range, 1)),

));

# NOTE: ind.downRank might be a hack
ind.downRankFull := (self, inds) >> inds[self.n];
ind.downRank := (self, loopid, ind) >> Cond(loopid=self.n, ind, self);
ind.rank := self >> self.n;

_hofnew := true;
# ExpMarkActiveRank(<s>)
#   performs a recursive walk over <s> and sets ._expMarkActiveRank attribute to the "active rank"
#   of each subexpression.
#
#   Active rank of an expression denotes the maximum implicit loop id that the expression refers to.
#   Implicit loop id's are introduced by objects such as GT and Lambda.
#
ExpMarkActiveRank := function(s)
    local c, rch, rank;
    rch := Cond(IsRec(s) and IsBound(s.rChildren), s.rChildren(), IsList(s) and not IsString(s), s, []);
    if ObjId(s) = ind then
        rank := s.n;
    else
	if _hofnew then
	    rank := _rank(s);
	    DoForAll(rch, ExpMarkActiveRank);
	else
# this was invalid (fcurry, fcall, etc) ->
	    rank := Maximum(_rank(s), Maximum0(List(rch, ExpMarkActiveRank)));
	fi;
    fi;

    if IsRec(s) and IsSymbolic(s) then s._expMarkActiveRank := rank; fi;
    return rank;
end;

_ExpMarkPassiveRank := function(s, parent_rank)
    local c, rch, rank, my_rank;
    rch := Cond(IsRec(s) and IsBound(s.rChildren), s.rChildren(), IsList(s) and not IsString(s), s, []);
    my_rank := _rank(s);

    # NOTE: this is a terrible hack! WHAT ABOUT Lambda?
    if spiral.paradigms.common.IsGT(s)        then rank := my_rank + parent_rank;
    elif ObjId(s)=ind then rank := Maximum(parent_rank, s.n);
    else                   rank := Maximum(parent_rank, my_rank);
    fi;
    if IsRec(s) and (IsSymbolic(s) or IsFuncExp(s)) then s._expMarkPassiveRank := rank; fi;

    DoForAll(rch, x -> _ExpMarkPassiveRank(x, rank));
    return rank;
end;
# ExpMarkPassiveRank(<s>)
#   performs a recursive walk over <s> and sets ._expMarkPassiveRank attribute to the "passive rank"
#   of each subexpression.
#
#   Passive rank of an expression denotes the maximum implicit loop id that is defined in the expression.
#   Regardless of whether expression refers to it or now. In contracst, an "active rank" is the loop id
#   that is actually referred to.
#
#   Implicit loop id's are introduced by objects such as GT and Lambda.
#
ExpMarkPassiveRank := s -> _ExpMarkPassiveRank(s, 0);

Class(lambdaWrap, Exp, rec(
    computeType := self >> Checked(ObjId(self.args[1].t) = TFunc, Last(self.args[1].t.params))
));


# Copyright (c) 2018-2021, Carnegie Mellon University
# See LICENSE for details

#F threadId()
#F
Class(threadId, Exp, rec(
    computeType := self >> TInt
));

#F barrier(<nthreads>, <tid>, <barrier-name>)
#F
#F Example of mapping from SMPBarrier object <o>:
#F    barrier(o.nthreads, o.tid, "&GLOBAL_BARRIER")
#F
Class(barrier, call, rec(
    visitAs := call
));

#F smp_fork(<nthreads>, <cmd>)
#F NOTE: this should take in thread id 'tid' for consistency, but it doesn't!
#F
Class(smp_fork, Command, rec(
    rChildren := self >> [self.nthreads, self.cmd],
    rSetChild := rSetChildFields("nthreads", "cmd"),

    __call__ := (self, nthreads, cmd) >> WithBases(self,
       rec(operations := CmdOps,
           nthreads   := Checked(IsInt(nthreads) or IsScalar(nthreads), nthreads),
           cmd        := Checked(IsCommand(cmd), cmd))),

    print := (self,i,is) >> Print(
         self.__name__, "(", self.nthreads, ",\n",
         Blanks(i+is), self.cmd.print(i+is, is), "\n",
         Blanks(i), ")"
    )
));

#F smp_loop(<nthreads>, <tidvar>, <tidexp>, <loopvar>, <range>, <cmd>)
#F
#F nthreads - # of threads
#F tidvar - thread id variable (will be set to tid value), necessary for nested parallelism
#F tidexp - thread id value
#F
Class(smp_loop, loop_base, rec(
   __call__ := meth(self, nthreads, tidvar, tidexp, loopvar, range, cmd)
       local result;
       Constraint(IsVar(loopvar));
       Constraint(IsCommand(cmd));
       range := toRange(range);
       if range = 0 then
           return skip();
       else
           loopvar.setRange(range);
           loopvar.isLoopIndex := true;
           return WithBases(self, rec(
                   operations := CmdOps,
                   nthreads := nthreads,
                   cmd := cmd,
                   var := loopvar,
                   tidvar := Checked(IsLoc(tidvar), tidvar),
                   tidexp := toExpArg(tidexp),
                   range := range));
       fi;
   end,

   rChildren := self >> [self.nthreads, self.tidvar, self.tidexp, self.var, self.range, self.cmd],
   rSetChild := rSetChildFields("nthreads", "tidvar", "tidexp", "var", "range", "cmd"),

   print := (self, i, is) >> Print(self.name, "(", self.nthreads, ", ",
       self.tidvar, ", ", self.tidexp, ", ", self.var, ", ",
       self.range, ",\n", Blanks(i+is),
       self.cmd.print(i+is, is),
       Print("\n", Blanks(i), ")")),

   DContainer := (self, o, y, x, opts) >> self(o.child(1), y, x, opts),

   free := self >> Difference(self.cmd.free(), [self.var, self.tidvar])
));

# Class(smp_chain, chain, rec(
#    __call__ := meth(arg)
#        local self, nthreads, cmds;
#        [self, nthreads, cmds] := [arg[1], arg[2], Flat(Drop(arg, 2))];
#        return WithBases(self, rec(
#                nthreads   := nthreads,
#                operations := CmdOps,
#                cmds       := Checked(ForAll(cmds, IsCommand), cmds)));
#    end,

#    print := (self,i,is) >> When(Length(self.cmds)=0,
#        Print(self.name, "(", self.nthreads, ")"),
#        Print(self.name, "(", self.nthreads, ",\n", self.printCmds(i+is, is), Blanks(i), ")"))
# ));

Class(SMPCodegenMixin, Codegen, rec(
    SMPBarrier := (self, o, y, x, opts) >> chain(
        self(o.child(1), y, x, opts),
        barrier(o.nthreads, o.tid, "&GLOBAL_BARRIER")),

    SMPSum := (self, o, y, x, opts) >> let(
        outer_tid     := When(IsBound(opts._tid), opts._tid, 0),
        outer_num_thr := When(IsBound(opts._tid), opts._tid.range, 1),
        tid := var.fresh("tid", TInt, o.nthreads * outer_num_thr),
        smp_loop(o.nthreads, tid, (outer_tid * outer_num_thr) + o.tid,
                 o.var, o.domain,
                 self(o.child(1), y, x, CopyFields(opts, rec(_tid := tid))))
    )
));

_PullBuffersSMP := function(expr, type_predicate, nthreads, tid)
    local ch, i, t, data, pullv, stayv;
    if ObjId(expr)=assign then return [[], expr];
    else
        # Implemented as a recursive tree walk
	data := [];

        if ObjId(expr)=decl then
            [pullv, stayv] := SplitBy(expr.vars, x->type_predicate(x.t));
            data := List(pullv, x->[x, nthreads, tid]);
            expr := decl(stayv, expr.cmd);
        elif ObjId(expr)=smp_loop then
            nthreads := nthreads * expr.nthreads;
            tid := expr.tidvar;
        fi;

        ch := _children(expr);
        for i in [1..Length(ch)] do
            t := _PullBuffersSMP(ch[i], type_predicate, nthreads, tid);
            Append(data, t[1]);
            _setChild(expr, i, t[2]);
        od;
        return [data, expr];
    fi;
end;

#F PullBuffersSMP(code, type_predicate)
#F
#F  Pulls out declarations (from decl(..)) that satisfy type_predicate,
#F  and return declared variables as triplets [var, nthreads, tid]
#F
#F  This function handled nested parallelism, thats why the implementation is
#F  not so straightforward.
#F
PullBuffersSMP := (code, type_predicate) -> _PullBuffersSMP(code, type_predicate, 1, 0);