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

#F ArgsExp(<exp>)
#F   Collects all location leaves in an expression.
#F   This is equivalent to Collect(exp, @.cond(IsLoc)), but has less overhead.
#F
ArgsExp := exp -> Cond(
    not IsRec(exp) or not IsSymbolic(exp), [],
    IsLoc(exp), Concatenation([exp], ConcatList(exp.rChildren(), ArgsExp)),
    IsValue(exp) or not IsBound(exp.args), [],
    Concatenation(List(exp.args, ArgsExp)));

assign.op_in  := self >> ConcatList(self.loc.rChildren(), ArgsExp) :: ArgsExp(self.exp);
assign.op_out := self >> [self.loc];
assign.op_inout := self >> [];

IsCmdOp := x -> IsCommand(x) and IsBound(x.op_in) and IsBound(x.op_out) and IsBound(x.op_inout);

VarArgsExp := exp -> Cond(
    IsVar(exp), [exp],
    IsValue(exp) or not IsRec(exp), [],
    ConcatList(exp.rChildren(), VarArgsExp));

#
# Def/Succ/Pred
#

MarkDefLoc := function(loc, d)
    loc.def := d;
    if not IsBound(loc.succ) then loc.succ := Set([]); fi;
end;

MarkPredLoc := function(loc, pred)
    pred := Set(ShallowCopy(pred));
    if not IsBound(loc.pred) then loc.pred := pred;
    # NOTE: below statement crashes with "<loc.pred> is not a set" error
    #        if UniteSet(loc.pred, pred) is used. This is an indication of
    #        broken < or = operations. Union() fixes the problem, but is
    #        essentially a hack
    else loc.pred := Union(loc.pred, pred);
    fi;

    if IsBound(loc.def) then
	if not IsBound(loc.def.pred) then loc.def.pred := pred;
	else loc.def.pred := Union(loc.def.pred, pred);
	fi;
    fi;
end;

# succesors is NOT a set to allow multiplicity (2 uses from same statement, like a = b*b, succ(b)=[a, a])
# this aids compiler from preventing propagation of b in this case in CopyPropagate.propagate()
#
MarkSuccLoc := function(loc, cmd_loc)
    if IsVar(loc) then
        if not IsBound(loc.succ) then loc.succ := [cmd_loc];
        else                          Add(loc.succ, cmd_loc); fi;

        if IsBound(loc.def) then
            if not IsBound(loc.def.succ) then loc.def.succ := [cmd_loc];
            else Add(loc.def.succ, cmd_loc);
            fi;
        fi;
    fi;
end;

DefLoc   := loc -> When(IsBound(loc.def),  loc.def,  false);
PredLoc  := loc -> When(IsBound(loc.pred), loc.pred, Set([]));
SuccLoc  := loc -> When(IsBound(loc.succ), loc.succ, ([]));

PredCmd  := cmd -> When(IsBound(cmd.pred), cmd.pred, Set([]));
SuccCmd  := cmd -> When(IsBound(cmd.succ), cmd.succ, ([]));

Declare(MarkDefUse, MarkPreds, ClearDefUse);

_ChainMarkDefUse := function(code)
    local c, u, pred;
    for c in code.cmds do
        if IsCmdOp(c) then
	    DoForAll(c.op_out(), x->MarkDefLoc(x, c));
	    pred := c.op_in() :: c.op_inout();
	    DoForAll(c.op_out(), x->MarkPredLoc(x, pred));
	    for u in pred do
	        DoForAll(c.op_out(), x->MarkSuccLoc(u, x));
	    od;
	else
	    DoForAll(c.rChildren(), MarkDefUse);
	fi;
    od;
    return code;
end;

_MarkDefUse := function(code)
    if not IsCommand(code) then return;
    elif IsChain(code) then _ChainMarkDefUse(code);
    else DoForAll(code.rChildren(), _MarkDefUse);
    fi;
    return code;
end;

MarkDefUse := function(code)
    code := ClearDefUse(code);
    return _MarkDefUse(code);
end;

_ChainMarkPreds := function(code)
    local c, u, pred;
    for c in code.cmds do
        if IsAssign(c) then
	    pred := c.op_in() :: c.op_inout();
	    MarkPredLoc(c.loc, pred);
	else
	    DoForAll(c.rChildren(), MarkPreds);
	fi;
    od;
    return code;
end;

_MarkPreds := function(code)
    if not IsCommand(code) then return;
    elif IsChain(code) then _ChainMarkPreds(code);
    else DoForAll(code.rChildren(), _MarkPreds);
    fi;
    return code;
end;

MarkPreds := function(code)
    code := ClearDefUse(code);
    return _MarkPreds(code);
end;

ClearDefUseLoc := function(loc)
   if IsBound(loc.def) then
       Unbind(loc.def.pred);
       Unbind(loc.def.succ);
   fi;
   Unbind(loc.def);
   Unbind(loc.succ);
   Unbind(loc.pred);
   return loc;
end;

depthLoc := x -> When(IsBound(x.depth), x.depth, 0);
rdepthLoc := (x, max) -> When(IsBound(x.rdepth), x.rdepth, max);

ComputeDepthsChain := function(code)
    local cmd, depth, rdepth, maxdepth, succs, dd, preds, succs;
    Constraint(ObjId(code)=chain);
    for cmd in code.cmds do
        Constraint(IsAssign(cmd));
	preds := PredCmd(cmd);
	dd := List(preds, depthLoc);
        depth := When(dd=[], 0, 1 + Maximum(dd));
	cmd.allpreds := Union(Filtered(preds,x->IsVar(x) and IsBound(x.def)),
	    Union(List(preds, x->When(IsVar(x) and IsBound(x.def) and
			              IsBound(x.def.allpreds), x.def.allpreds, []))));
        cmd.earliest := depth;
	cmd.loc.depth := depth;
    od;
    maxdepth := Length(code.cmds);
    for cmd in Reversed(code.cmds) do
	succs := SuccCmd(cmd);
        dd := List(succs, s -> rdepthLoc(s, maxdepth));
        rdepth := When(dd=[], maxdepth,  -1 + Minimum(dd));
	cmd.allsuccs := Union(succs,
	    Union(List(succs, x->When(IsBound(x.def) and
			              IsBound(x.def.allsuccs), x.def.allsuccs, []))));

        cmd.latest := rdepth;
	cmd.loc.rdepth := rdepth;

    od;
end;

#F ClearDefUse(<code)> - clears attributes set by MarkDefUse
#F
ClearDefUse := code -> Chain(DoForAll(Collect(code, var), ClearDefUseLoc), code);

DFSChain := function(code)
    local c, cmds, W, next, p, added, schedule;
    cmds := code.cmds; W := []; schedule := [];
    for c in cmds do
       if not IsBound(c.loc.generated) then
       Add(schedule, c); c.loc.generated := true;
       if IsBound(c.loc.succ) then
           Append(W, c.loc.succ);
       fi;
       fi;
       while Length(W)<>0 do
           next := Last(W);
       added := false;
       if IsBound(next.pred) then
           for p in next.pred do
               if not IsBound(p.generated) then Add(W, p); added := true; fi;
           od;
       fi;
       if not added then
           if IsBound(next.def) then Add(schedule, next.def); fi;
           next.generated := true;
           RemoveLast(W, 1);
       fi;
       od;
    od;

    for c in cmds do Unbind(c.loc.generated); od;
    return chain(schedule);
end;

InputsChain := code -> List(Filtered(code.cmds, x->not ForAny(x.loc.pred,p->IsBound(p.def))),x->x.loc);
OutputsChain := code -> List(Filtered(code.cmds, x->not IsBound(x.loc.succ) or x.loc.succ=[]), x->x.loc);

_ClearRedBlue := function(v) Unbind(v.blue); Unbind(v.red); end;

ClearRedBlue := code -> DoForAll(Collect(code, var), _ClearRedBlue);

_HSplitChain := function(c,inp,out)
    local i,o,x,inpsucc,outpred, mid,done,red,blue,rcmds,bcmds;
    mid := Set([]);
    blue := Set(inp); bcmds:=List(blue, x->x.def);
    red := Set(out); rcmds:=List(red, x->x.def);

    done := (inp=[]) and (out=[]);
    while not done do
        inp := Difference(Union(List(inp, x->x.succ)),blue); # unvisited successors
    out := Difference(Union(List(out, x->x.pred)),red); # unvisited predecessors

    for x in inp do
        AddSet(blue, x);
        if not (x in red) then Add(bcmds,x.def); fi;
    od;

    for x in out do
        AddSet(red, x);
        if not (x in blue) then Add(rcmds,x.def); fi;
    od;

    inp := Difference(inp, red);
    out := Difference(out, blue);

    done := (inp=[]) and (out=[]);
    od;
    return [Intersection(red,blue), bcmds, Reversed(rcmds)];
end;

HSplitChain := function(c)
    local inp, out, mid;
    inp := Set(InputsChain(c));
    out := Set(OutputsChain(c));
    mid := _HSplitChain(c, inp, out);
    return mid;
end;

scheduled := code -> Inherit(code, rec(cmd := DFSChain(code.cmd)));


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

Rows := s -> Cond(
    IsMat(s),
        Length(s),

    IsBound(s.rng) and IsBound(s.dmn),
        s.dims()[1],

    IsValue(s) or IsSymbolic(s),
        Checked(IsArrayT(s.t), IsArrayT(s.t.t),
	    s.t.size),

    s.dimensions[1]
);

Cols := s -> Cond(
    IsMat(s),
        Length(s[1]),

    IsBound(s.rng) and IsBound(s.dmn),
        s.dims()[2],

    IsValue(s) or IsSymbolic(s),
        Checked(IsArrayT(s.t), IsArrayT(s.t.t),
	    s.t.t.size),

    s.dimensions[2]
);

olRows := s -> Flat([Rows(s)]);
olCols := s -> Flat([Cols(s)]);

#F SPL(<rec>)
#F    Set operations field to SPLOps. This function should be called
#F    to initialize SPL instances.
SPL := function(record)
   record.operations := SPLOps;
   return record;
end;

HashAsSPL := o -> Cond(
    IsList(o),
        List(o, HashAsSPL),
    not IsRec(o) or (not IsBound(o.hashAs) and not IsBound(o.from_rChildren)),
        o,
    IsBound(o.hashAs),
        o.hashAs(),
    IsValue(o),
        o.v,
    o.from_rChildren(List(o.rChildren(), HashAsSPL))
);

# ==========================================================================
# ClassSPL
#
# Base class for all SPL constructs
# ==========================================================================
Class(ClassSPL, AttrMixin, rec(
    isSPL := true,
    transposed := false,

#DD this allows objects which are not TaggedNonTerminals to drop tags automagically
#DD and without error.
    withTags := (self, t) >> self,

    _short_print := false,
    _newline := i ->  Print("\n", Blanks(i)),
    _indent  := i ->  Print(Blanks(i)),
    _indentStr := Blanks,

    __call__ := meth(arg)
        local self, params, nump, A,p,res,h,lkup;
        self := arg[1];
        params := arg{[2..Length(arg)]};
        nump := Length(params);

        if not IsBound(self.new)  then
            Error("Constructor for this class is not implemented");
        elif IsBound(self.abbrevs) and self.abbrevs <> [] then
            for A in self.abbrevs do
                if NumArgs(A) = -1 or NumArgs(A) = nump then
                    params := ApplyFunc(A, params);
                fi;
            od;
        fi;

        if NumArgs(self.new)-1 <> Length(params) then
            Error("Constructor requires ", NumArgs(self.new)-1, " parameters (",
                Length(arg)-1, " given): ",
                ParamsMeth(self.new));
        else
            h := self.hash;
            if h<>false then
                lkup := h.objLookup(self, params);
                if lkup[1] <> false then return lkup[1]; fi;
            fi;

            res := ApplyFunc(self.new, params);

            if h<>false then return h.objAdd(res, lkup[2]);
            else return res;
            fi;
        fi;
    end,

    hash := false,

    checkDims := self >> DimensionsMat(MatSPL(self)) = self.dimensions,

    #-----------------------------------------------------------------------
    # create a new object with .<name> field set to true
    setAttr := meth(self, name)
        local s;
        s:= Copy(self);
        s.(name) := true;
        return s;
    end,
    # ----------------------------------------------------------------------
    # create a new object with .<name> field set to <val>
    setAttrTo := meth(self, name, val)
        local s;
        s:= Copy(self);
        s.(name) := val;
        return s;
    end,

    #---------Backwards Compatibility for the new dimension system------
    setDims := meth(self) self.dimensions := self.dims(); return self; end,

    dims := self >> [ StripList(List(self.rng(), l -> l.size)),
	              StripList(List(self.dmn(), l -> l.size)) ],

    advdims := (self) >> let(d := self.dims(), [ [[ d[1] ]], [[ d[2] ]] ]),
    arity   := (self) >> List(self.dims(), e -> Length(Flat([e]))),

    TType:=TUnknown,

    rng := meth(self) local d;
        if IsBound(self.dims) then
            d := Flat([self.dims()[1]]);
        else
            d := [self.dimensions[1]];
        fi;
        if IsBound(self.a.t_out) then
            return List(TransposedMat([self.a.t_out, d]), e -> TArray(e[1], e[2]));
        else
            return List(d, e -> TArray(self.TType, e));
        fi;
    end,

    dmn := meth(self) local d;
        if IsBound(self.dims) then
            d := Flat([self.dims()[2]]);
        else
            d := [self.dimensions[2]];
        fi;
        if IsBound(self.a.t_in) then
            return List(TransposedMat([self.a.t_in, d]), e -> TArray(e[1], e[2]));
        else
            return List(d, e -> TArray(self.TType, e));
        fi;
    end,

    free := self >> Union(List(self.rChildren(), FreeVars)),

    equals := (self, o) >>
        ObjId(self) = ObjId(o) and self.rChildren() = o.rChildren() and self.a = o.a,

    lessThan := (self, o) >> Cond(
        ObjId(self) <> ObjId(o), ObjId(self) < ObjId(o),
        [ ObjId(self), self.rChildren(), self.a ] < [ ObjId(o), o.rChildren(), o.a ]
    ),

    from_rChildren := (self, rch) >> ApplyFunc(ObjId(self), rch).appendAobj(self),

    print := (self, i, is) >> self._print(self.rChildren(), i, is),

    _print := meth(self, ch, indent, indentStep)
        local s, first, newline;

	if self._short_print or ForAll(ch, x->not IsRec(x) or IsSPLSym(x) or IsSPLMat(x)) then
	    newline := Ignore;
	else
	    newline := self._newline;
	fi;

	first := true;
        Print(self.__name__, "(");
	for s in ch do
            if(first) then first:=false;
            else Print(", "); fi;
            newline(indent + indentStep);
            When(IsSPL(s) or (IsRec(s) and IsBound(s.print) and NumGenArgs(s.print)=2),
                 s.print(indent + indentStep, indentStep),
		 Print(s));
	od;
	newline(indent);
	Print(")");
        self.printA();

	if IsBound(self._setDims) then
            Print(".overrideDims(", self._setDims, ")");
	fi;
    end,

    printlatex := meth(self)
        local s, first, newline, i;

	first := true;
        Print("(");
        i := Length(self.rChildren());
        for s in self.rChildren() do
            When(IsSPL(s) or (IsRec(s) and IsBound(s.printlatex) and NumGenArgs(s.printlatex)=0),
                 s.printlatex(),
		 Print(s));
                 if i >=2 then
                    Print(" ", When(IsBound(self.latexSymbol), self.latexSymbol, ""), " ");
                 fi;
                 i := i-1;
	od;
	Print(")");
    end,

    overrideDims := (self, dims) >> CopyFields(self, rec(_setDims := dims, dimensions := dims)),

    terminate     := self >> self.from_rChildren(List(self.rChildren(), x->When(IsSPL(x), x.terminate(), x))),

    # ------------------------- Required methods ---------------------------\
    transposeSymmetric := True,
    dims          := meth(self) Error("Not implemented"); end,
    isPermutation := meth(self) Error("Not implemented"); end,
    isTerminal    := meth(self) Error("Not implemented"); end,
    isReal        := meth(self) Error("Not implemented"); end,
    isInplace     := self >> Rows(self)=Cols(self) and let(ch:=self.children(), Cond(Length(ch)=0, false, ForAll(ch, x->x.isInplace()))),
    children      := meth(self) return []; end,
    numChildren   := meth(self) return 0; end,
    child         := meth(self,n) Error("Not implemented"); end,
    setChild      := meth(self,n,what) Error("Not implemented"); end,
    toAMat        := meth(self) Error("Not implemented"); end,
    transpose     := meth(self) Error("Not implemented"); end,
    conjTranspose  := meth(self) Error("Not implemented"); end,
));

#F <SPL> * <SPL>
#F <scalar> * <SPL>
#F   is equivalent to ComposeSPL and ScalarMultiple resp.
#F
SPLOps.\* := (S1, S2) ->
    Cond(IsSPL(S1) and IsSPL(S2),   Compose(S1, S2),
         IsSPL(S2),                 Scale(S1, S2),
     IsSPL(S1),                 Scale(S2, S1),
     Error("do not know how to compute <S1> * <S2>"));

#F <SPL> + <SPL>
#F   is equivalent to SUM(<S1>, <S2>)
#F
SPLOps.\+ := (S1, S2) ->
    Cond(IsSPL(S1) and IsSPL(S2),   SUM(S1, S2),
     Error("do not know how to compute <S1> + <S2>"));

#F S1 ^ S2
#F   is equivalent to ConjugateSPL(S1, S2).
#F
SPLOps.\^ := (S1, S2) ->
    When(IsSPL(S1) and IsSPL(S2),
         Conjugate(S1, S2),
     Error("do not know how to compute S1 ^ S2"));


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

# utility functions

Range0 := (a) -> When(IsList(a), [0..Length(a)-1], [0..a-1]);
Range1 := (a) -> When(IsList(a), [1..Length(a)], [1..a]);

# special rules for annotating the recursion level for all variables.
# allows us to just keep the variables, instead of also needing to keep
# the ISums.

_checkTag := function(e, cx)
    local sums;

    if IsBound(e.var.order) then
        return false;
    fi;

    sums :=  Filtered(cx.parents, e -> ObjId(e) = ISum);

    return (sums = [] or IsBound(Last(sums).var.order));
end;

_setTag := function(e, cx)
    local sums;

    sums :=  Filtered(cx.parents, e -> ObjId(e) = ISum);

    if sums = [] then
        e.var.order := 0;
    else
        e.var.order := Last(sums).var.order + 1;
    fi;

    return e;
end;

Class(_MissEstPreprocessJams, RuleSet);
Class(_MissEstSanitize, RuleSet);
Class(_MissEstRules, RuleSet);
Class(_MissEstCleanup, RuleSet);
Class(_MissEstInplace, RuleSet);

Class(_GSWrap, SumsBase, rec(
    __call__ := (self, rng, dmn, payload) >> WithBases(self, rec(
        _dmn := dmn,
        _rng := rng,
        _payload := payload,
        _children := []
    )),
    rChildren := (self) >> [],
    dims := (self) >> [
        StripList(List(self.rng(), (l) -> l.size)),
        StripList(List(self.dmn(), (l) -> l.size))
    ],
    rSetChild := (self, n, what) >> Error("no kids"),
    rng := (self) >> self._rng,
    dmn := (self) >> self._dmn,
    print := (self,i,is) >> Print(self.name)
));

Class(_AltBB, BB);
Class(_AltInplace, Inplace);

# wrapper needs an identifier, for later matching since two wraps
# are always spawned. See _MissEstInplace.InplaceExpand rule.
#
Class(_InplaceWrap, _GSWrap);
#SumsBase, BaseMat, rec(
#    new := (self, id) >> SPL(WithBases(self, rec(
#        id := id
#    )))
#));

Class(_KeepScat, Scat);
Class(_KeepGath, Gath);

RewriteRules(_MissEstPreprocessJams, rec(

 ComposeGathGath := ARule(Compose, [ @(1, Gath), @(2, [Gath, Prm]) ], # o 1-> 2->
     e -> [ Gath(fCompose(@(2).val.func, @(1).val.func)) ]),

 ComposeScatScat := ARule(Compose, [ @(1, [Scat, ScatAcc]), @(2, [Scat, ScatAcc]) ], # <-1 <-2 o
     e -> [ Cond(ObjId(@(1).val)=ScatAcc or ObjId(@(2).val)=ScatAcc,
                 ScatAcc(fCompose(@(1).val.func, @(2).val.func)),
                 Scat   (fCompose(@(1).val.func, @(2).val.func))) ]),

 PullInRight := ARule( Compose,
       [ @(1, [Prm, Scat, ScatAcc, TCast, PushR, PushLR, Conj, ConjL, ConjR, ConjLR, FormatPrm]),
         @(2, [RecursStep, Grp, BB, SUM, JamISum, Buf, ISum, ICompose, Data, COND, NoDiagPullin, NoDiagPullinLeft, NoDiagPullinRight, NeedInterleavedComplex]) ],
  e -> [ CopyFields(@(2).val, rec(
             _children :=  List(@(2).val._children, c -> @(1).val * c),
             dimensions := [Rows(@(1).val), Cols(@(2).val)] )) ]),

 PullInLeft := ARule( Compose,
       [ @(1, [RecursStep, Grp, BB, SUM, SUMAcc, JamISum, Buf, ISum, ICompose, ISumAcc, Data, COND, NoDiagPullin, NoDiagPullinLeft, NoDiagPullinRight, NeedInterleavedComplex]),
         @(2, [Prm, Gath, TCast, PushL, PushLR, Conj, ConjL, ConjR, ConjLR, FormatPrm]) ],
     e -> [ CopyFields(@(1).val, rec(
                _children := List(@(1).val._children, c -> c * @(2).val),
                dimensions := [Rows(@(1).val), Cols(@(2).val)] )) ]),
));

RewriteRules(_MissEstSanitize, rec(

    # sometimes there is a BB inside a BB. this is not good for me.
    # this removes all interior BBs
    RemoveInnerBB := Rule(@@(1,BB,(e,cx) -> ForAny(cx.parents, e -> ObjId(e) = BB)), e -> @@(1).val.rChildren()[1]),

    # BB cannot have an inplace inside of it, since all operations inside
    # the BB do not spill to memory.
    RemoveInnerInplaceFromBB := Rule(@@(1,Inplace,(e,cx) -> ForAny(cx.parents, e -> ObjId(e) = BB)), e -> @@(1).val.rChildren()[1]),

    # same thing with Inplace wrappers. leave only the outer.
    RemoveInnerInplace := Rule(@@(1,Inplace,(e,cx) -> ForAny(cx.parents, e -> ObjId(e) = Inplace)), e -> @@(1).val.rChildren()[1])
));

# we tag the BB with an inplace tag, if that basic block is to be performed
# inplace.

_tagBBin := function(bb, v)
    bb.inplin := v;
    return bb;
end;

_tagBBout := function(bb, v)
    bb.inplout := v;
    return bb;
end;

_tagInplace := function(k, v)
    k.inplace := v;
    return k;
end;

#
# this thing works very similarly to the keepgath/keepscat below. we start with
# the inplace tag, generate some wrappers from it, propagate the wrappers
# into the sigm-spl tree looking for a BB. once the wrappers hit the BB, they
# tag it.

RewriteRules(_MissEstInplace, rec(
    InplaceExpand := Rule(Inplace,  e ->
        let(i := Ind(2),
            _AltInplace(Compose(Concat(
                [_InplaceWrap(e.rng(), e.rng(), i)],
                e.rChildren(),
                [_InplaceWrap(e.dmn(), e.dmn(), i)]
            )))
        )
    ),

    WrapCompose := ARule(Compose, [@(1, _InplaceWrap), @(2, Compose)], e -> [
        Compose(
            Concat([@(1).val], @(2).val.rChildren())
        )
    ]),

    ComposeWrap := ARule(Compose, [@(1, Compose), @(2, _InplaceWrap)], e -> [
        Compose(
            Concat(@(2).val.rChildren(), [@(1).val])
        )
    ]),

    WrapISum := ARule(Compose, [@(2, _InplaceWrap), @(1, [ISum, JamISum])], e -> [
        CopyFields(@(1).val, rec(
            _children := [
                Compose(
                    @(2).val, # _GSWrap(@(1).val.rng(), @(1).val.rng()),
                    @(1).val.child(1)
                )
            ],
            dimensions := [Rows(@(1).val), Cols(@(2).val)]
        ))
    ]),

    ISumWrap := ARule(Compose, [@(1, [ISum, JamISum]), @(2, _InplaceWrap)], e -> [
        CopyFields(@(1).val, rec(
            _children := [
                Compose(
                    @(1).val.child(1),
                    @(2).val # _GSWrap(@(1).val.rng(), @(1).val.dmn()) # have to adjust the size as we move it in.
                )
            ],
            dimensions := [Rows(@(1).val), Cols(@(2).val)]
        ))
    ]),

    InplaceWrapBB := ARule(Compose, [@(1, _InplaceWrap), @(2,BB)], e -> [
        _tagBBout(@(2).val, @(1).val._payload)
    ]),

    BBInplaceWrap := ARule(Compose, [@(1, BB), @(2, _InplaceWrap)], e -> [
        _tagBBin(@(1).val, @(2).val._payload)
    ]),

    ComposeAssoc := ARule( Compose, [ @(1,Compose) ],  e -> @(1).val.children() )

));

#
# we push into the basic block to figure out the gath/scat which access memory.
#
RewriteRules(_MissEstRules, rec(

    # insert a first gather and last scatter. replace BB with
    # a fake for a while.
    BBExpand := Rule(BB,  e ->
        _AltBB(Compose(Concat(
            [_GSWrap(e.rng(), e.rng(), When(IsBound(e.inplout), e.inplout, false))],
            e.rChildren(),
            [_GSWrap(e.dmn(), e.dmn(), When(IsBound(e.inplin), e.inplin, false))])
        ))
    ),
    WrapCompose := ARule(Compose, [@(1, _GSWrap), @(2, Compose)], e -> [
        Compose(
            Concat([@(1).val], @(2).val.rChildren())
        )
    ]),

    ComposeWrap := ARule(Compose, [@(1, Compose), @(2, _GSWrap)], e -> [
        Compose(
            Concat(@(2).val.rChildren(), [@(1).val])
        )
    ]),

    WrapISum := ARule(Compose, [@(2, _GSWrap), @(1, [JamISum, ISum])], e -> [
        CopyFields(@(1).val, rec(
            _children := [
                Compose(
                    @(2).val, # _GSWrap(@(1).val.rng(), @(1).val.rng()),
                    @(1).val.child(1)
                )
            ],
            dimensions := [Rows(@(1).val), Cols(@(2).val)]
        ))
    ]),

    ISumWrap := ARule(Compose, [@(1, [JamISum, ISum]), @(2, _GSWrap)], e -> [
        CopyFields(@(1).val, rec(
            _children := [
                Compose(
                    @(1).val.child(1),
                    @(2).val # _GSWrap(@(1).val.rng(), @(1).val.dmn()) # have to adjust the size as we move it in.
                )
            ],
            dimensions := [Rows(@(1).val), Cols(@(2).val)]
        ))
    ]),

    GSWrapScat := ARule(Compose, [@(1, _GSWrap), @(2,Scat)], e -> [
        _tagInplace(_KeepScat(@(2).val.func), @(1).val._payload)
    ]),

    GathGSWrap := ARule(Compose, [@(1, Gath), @(2, _GSWrap)], e -> [
        _tagInplace(_KeepGath(@(1).val.func), @(2).val._payload)
    ]),

    ComposeAssoc := ARule( Compose, [ @(1,Compose) ],  e -> @(1).val.children() )
));

RewriteRules(_MissEstCleanup, rec(
    ISumTagVar := Rule(@@(1,[JamISum, ISum], _checkTag), _setTag),

    RemoveBlk := ARule(Compose, [Blk], e -> []),
    RemoveRCDiag := ARule(Compose, [RCDiag], e -> []),

#    RemoveScatGath := ARule(Compose, [Scat, Gath], e -> []),
#    RemoveGath := ARule(Compose, [Gath], e -> []),
#    RemoveScat := ARule(Compose, [Scat], e -> []),
    # collapse ISums with just one G or S inside.
#    ISumGath := Rule([ISum, @(1, [Gath, Scat])], e -> @(1).val),

    ComposeAssoc := ARule( Compose, [ @(1,Compose) ],  e -> @(1).val.children() ),

    AltBBtoBB := Rule(@(1, _AltBB), e -> BB(@(1).val.child(1))),

    AltInplacetoInplace := Rule(@(1,_AltInplace), e -> Inplace(@(1).val.child(1))),

#    ComposeAltInplacetoInplace := ARule(Compose, [@(1, _AltInplace)], e ->
#        [Inplace(@(1).val.child(1))]
#    ),

#    buggy rule. huh?
#    DropInplaceWrap := ARule(Compose, [@(1, _InplaceWrap)], e -> [])
));

# plot 1
# naive iterative inplace WHT(1M)

_returnAndReset := function(a)
    local t;

    t := a.switch;

    a.switch := false;

    return t;
end;

_set := function(a, v)
    a.switch := v;
end;

Class(cachesim, rec(
    __call__ := (self, e, s, a) >> WithBases(self, rec(
        e := e,
        s := s,
        a := a,

        tagstore := List(Range1(s), e ->
            List(Range1(a), ee ->
                rec(tag:=-1, time:=0)
            )
        ),
        count := 0,
		hash := false
    )),

    reset := meth(self)
        self.tagstore := List(Range1(self.s), e ->
            List(Range1(self.a), ee ->
                rec(tag:=-1, time:=0)
            )
        );
        self.count := 0;
    end,

    access := meth(self, addr)
        local setidx, tag, tagidx, new, bits, mask, shift;

		if self.hash then
			bits := Log2Int(self.a * self.s);
			mask := self.e * ((self.a * self.s) - 1);
			shift := bits;

			# we clip it (somewhat arbitrarily) at 30 bits.
			while (shift < 30) do
				addr := BinXor(
					addr,
					BinAnd(
						mask,
						Int(addr / shift)
					)
				);
				shift := shift + bits;
			od;
		fi;

		# add 1 for array offset
       	setidx := 1 + (Int(addr / self.e) mod self.s);

        tag := Int(addr / (self.e * self.s));

        tagidx := Filtered(Range1(self.a), e ->
            tag = self.tagstore[setidx][e].tag
        );

        # paranoia
        Constraint(Length(tagidx) <= 1);

        if Length(tagidx) = 0 then
            tagidx := self.a;
            self.tagstore[setidx][tagidx].tag := tag;
            new := 1;
        else
            tagidx := tagidx[1];
            new := 0;
        fi;

        self.count := self.count + 1;
        self.tagstore[setidx][tagidx].time := self.count;

        # sort based on access time, promotion to MRU happens here.
        Sort(self.tagstore[setidx], (a,b) -> a.time > b.time);

		# AppendTo("accesses", addr, " (", setidx, ", ", tagidx, "): ", new, "\n");
        return new;
    end,

    evaljam := meth(self, l, spl, res, jams)
        local l2, last, i, memaddr;

        if jams = [] then
            # get the actual address.
            memaddr := spl.array.offset + l.eval().v;

			# AppendTo("accesses", When(ObjId(spl) = _KeepGath, "G ", "S "));
            # perform the memory access!
            res.activ := res.activ + self.access(memaddr);
            res.access := res.access + 1;
        else

            last := Last(jams);

            for i in Range0(last.range) do
                l2 := Copy(l);

                SubstVars(l2, tab((last.id) := V(i)));

                self.evaljam(l2, spl, res, DropLast(jams, 1));

            od;
        fi;
    end,

	evalinner := meth(self, l, spl, res, inners, jams)
		local i, l2, last;

		if inners = [] then
        	for i in Range0(l.vars[1].range) do

            	l2 := Copy(l.at(i));

            	# now we have to deal with multiple jammed loops
            	# that means we have a nice little recursion here.
            	self.evaljam(l2, spl, res, jams);
        	od;
		else
			last := Last(inners);

			for i in Range0(last.range) do
				l2 := Copy(l);

				SubstVars(l2, tab((last.id) := V(i)));

				self.evalinner(l2, spl, res, DropLast(inners, 1), jams);
			od;
		fi;
	end,

    evalspl := meth(self, spl, data)
        local res, id, i, j, tres, l, l2;
        res := rec(access := 0, activ := 0);

        id := ObjId(spl);

        if id = ISum then

			# we evaluate ISums inside of BB differently
			if data.currsize <= data.K then
				# AppendTo("traverse", "inner ", spl.var.id, " ", spl.var.range,", ", data.currsize, "\n");
				Add(data.innervar, spl.var);
				data.currsize := data.currsize / spl.var.range;
            	tres := self.evalspl(spl.child(1), data);
				data.currsize := data.currsize * spl.var.range;
            	data.innervar := DropLast(data.innervar, 1);

            	res.access := res.access + tres.access;
            	res.activ := res.activ + tres.activ;

			else
				# AppendTo("traverse", "outer", spl.var.id, " ", spl.var.range,", ", data.currsize, "\n");
				data.currsize := data.currsize / spl.var.range;
            	for i in Range0(spl.var.range) do
                	Add(data.var, spl.var.id);
                	Add(data.val, V(i));
                	tres := self.evalspl(spl.child(1), data);
                	data.var := DropLast(data.var, 1);
                	data.val := DropLast(data.val, 1);

                	res.access := res.access + tres.access;
                	res.activ := res.activ + tres.activ;
            	od;
				data.currsize := data.currsize * spl.var.range;
			fi;

		elif id = BB then
			data.inBB := true;
            tres := self.evalspl(spl.child(1), data);
			data.inBB := false;

            res.access := res.access + tres.access;
            res.activ := res.activ + tres.activ;

        elif id = JamISum then
            Add(data.jamvar, spl.var);
            tres := self.evalspl(spl.child(1), data);
            data.jamvar := DropLast(data.jamvar, 1);

            res.access := res.access + tres.access;
            res.activ := res.activ + tres.activ;

        elif id = Compose then
            for i in Reversed(spl.children()) do
                tres := self.evalspl(i, data);
                res.access := res.access + tres.access;
                res.activ := res.activ + tres.activ;
            od;

        elif id = _KeepGath or id = _KeepScat then

            l := spl.func.lambda();

            # subst the normal loops
            DoForAll(Range1(data.var), e ->
                SubstVars(l, tab((data.var[e]) := data.val[e]))
            );

			# handle the innermost outer loop, nice little
			# recursion here.
			self.evalinner(l, spl, res, data.innervar, data.jamvar);

        else
            tres := List(spl.children(), e -> self.evalspl(e, data));
            res.access := Sum(List(tres, e -> e.access));
            res.activ := Sum(List(tres, e -> e.activ));
        fi;

        return res;
    end,

	applyRules := (self, s, o) >>
        ApplyStrategy(s, [
            _MissEstPreprocessJams,
            RulesFuncSimp,
            _MissEstSanitize,
            _MissEstInplace,
            _MissEstRules,
            _MissEstCleanup
        ], UntilDone, o),

    eval := meth(self, s, _K)
        local gath, scat, gf, sf, arrays, i;

        gath := Reversed(Collect(s, _KeepGath));
        scat := Reversed(Collect(s, _KeepScat));

        arrays := [TArray(TUnknown, gath[1].func.range())];

		# PrintTo("accesses", "\n");
		# PrintTo("traverse", "\n");

        # assign input/output arrays
        for i in Range1(gath) do

            # paranoia
            # Constraint(gath[i].func.domain() = scat[i].func.domain());

            gath[i].array := Last(arrays);

			# we can match any of the previous gather matrices.
			if scat[i].inplace <> false
				and ForAny(gath{[1..i]}, e -> e.inplace = scat[i].inplace) then

				scat[i].array :=
					Filtered(gath{[1..i]}, e -> e.inplace = scat[i].inplace)[1].array;
			else

                Add(arrays, TArray(TUnknown, scat[i].func.range()));
                scat[i].array := Last(arrays);
            fi;
        od;

        # descending order by size of array
        Sort(arrays, (a,b) -> a.size > b.size);

        # mark offsets
        arrays[1].offset := 0;

        # propagate offsets
        for i in [2..Length(arrays)] do
            arrays[i].offset := arrays[i-1].offset + arrays[i-1].size;
        od;

        # traverse the spl expression.
        return self.evalspl(s, rec(
			m := s.dims()[1],
			K := _K,

			currsize := s.dims()[1],
			var:=[],
			val:=[],
			innervar := [],
			jamvar:=[]
		));
    end,
));