#############################################################################
####
##
#W  anupqopt.gi                ANUPQ package                    Werner Nickel
#W                                                                Greg Gamble
##
##  Install file for functions to do with option manipulation.
##
#Y  Copyright (C) 2001  Lehrstuhl D fuer Mathematik,  RWTH Aachen,  Germany
##

#############################################################################
##
#V  PQ_FUNCTION . . . . . . . . . internal functions called by user functions
##
##  A record whose fields are (function)  names  and  whose  values  are  the
##  internal functions called by the functions with those names.
##
InstallValue( PQ_FUNCTION,
              rec( Pq                   := PQ_EPI_OR_PCOVER,
                   PqDescendants        := PQ_DESCENDANTS,
                   StandardPresentation := PQ_EPIMORPHISM_STANDARD_PRESENTATION,
                   PqDescendantsTreeCoclassOne := PqDescendantsTreeCoclassOne
                  )
             );

#############################################################################
##
#V  ANUPQoptions  . . . . . . . . . . . . . . . . . . . .  admissible options
##
##  is a record of lists of names of admissible {\ANUPQ} options,  such  that
##  each field is either the name of a (``key'') {\ANUPQ}  function  and  the
##  corresponding value is the list of option names that are  admissible  for
##  the function.
##
InstallValue( ANUPQoptions,
              rec( # options for `Pq' and `PqEpimorphism'
                   Pq  := [ "Prime",
                            "ClassBound",
                            "Exponent",
                            "Metabelian",
                            "OutputLevel",
                            "Relators",
                            "Identities",
                            "GroupName",
                            "SetupFile",
                            "PqWorkspace",
                            "RedoPcp" ],

                   # options for `PqDescendants'
                   PqDescendants
                       := [ "ClassBound",
                            "OrderBound",
                            "Relators",
                            "GroupName",
                            "StepSize",
                            "PcgsAutomorphisms",
                            "RankInitialSegmentSubgroups",
                            "SpaceEfficient",
                            "CapableDescendants",
                            "AllDescendants",
                            "Exponent",
                            "Metabelian",
                            "SubList",
                            "BasicAlgorithm",
                            "CustomiseOutput",
                            "SetupFile",
                            "PqWorkspace" ],

                   # options for `[Epimorphism][Pq]StandardPresentation'
                   StandardPresentation
                       := [ "Prime",
                            "pQuotient",
                            "ClassBound",
                            "Relators",
                            "GroupName",
                            "PcgsAutomorphisms",
                            "Exponent",
                            "Metabelian",
                            "OutputLevel",
                            "StandardPresentationFile",
                            "SetupFile",
                            "PqWorkspace" ],

                   # options for `PqDescendantsTreeCoclassOne'
                   PqDescendantsTreeCoclassOne
                       := [ "ClassBound",
                            "OrderBound",
                            "Relators",
                            "GroupName",
                            "StepSize",
                            "PcgsAutomorphisms",
                            "RankInitialSegmentSubgroups",
                            "SpaceEfficient",
                            "CapableDescendants",
                            "AllDescendants",
                            "Exponent",
                            "Metabelian",
                            "SubList",
                            "BasicAlgorithm",
                            "CustomiseOutput",
                            "TreeDepth",
                            "SetupFile",
                            "PqWorkspace" ],

                   PqList
                       := [ "SubList" ],
                   PqPcPresentation
                       := [ "Prime",
                            "ClassBound",
                            "Exponent",
                            "Metabelian",
                            "OutputLevel",
                            "Relators",
                            "Identities",
                            "GroupName" ],
                   PqNextClass
                       := [ "QueueFactor" ],
                   PqEvaluateIdentities
                       := [ "Identities" ],
                   PqDoExponentChecks
                       := [ "Bounds" ],
                   PqDisplayStructure
                       := [ "Bounds" ],
                   PqDisplayAutomorphisms
                       := [ "Bounds" ],
                   PqSPComputePcpAndPCover
                       := [ "Prime",
                            "ClassBound",
                            "Exponent",
                            "Metabelian",
                            "OutputLevel",
                            "Relators",
                            "GroupName" ],
                   PqSPSavePresentation
                       := [ "ClassBound",
                            "PcgsAutomorphisms",
                            "StandardPresentationFile" ],
                   PqPGSetDescendantToPcp
                       := [ "Filename" ],
                   PqPGSupplyAutomorphisms
                       := [ "NumberOfSolubleAutomorphisms",
                            "RelativeOrders" ],
                   PqPGConstructDescendants
                       := [ "ClassBound",
                            "OrderBound",
                            "StepSize",
                            "PcgsAutomorphisms",
                            "RankInitialSegmentSubgroups",
                            "SpaceEfficient",
                            "CapableDescendants",
                            "AllDescendants",
                            "Exponent",
                            "Metabelian",
                            "BasicAlgorithm",
                            "CustomiseOutput" ],
                   PqAPGDegree
                       := [ "Exponent" ],
                   PqAPGPermutations
                       := [ "PcgsAutomorphisms",
                            "SpaceEfficient",
                            "PrintAutomorphisms",
                            "PrintPermutations" ],
                   PqAPGOrbits
                       := [ "PcgsAutomorphisms",
                            "SpaceEfficient",
                            "CustomiseOutput" ],
                   PqAPGOrbitRepresentatives
                       := [ "PcgsAutomorphisms",
                            "SpaceEfficient",
                            "CapableDescendants",
                            "AllDescendants",
                            "Exponent",
                            "Metabelian",
                            "CustomiseOutput",
                            "Filename" ],
                   PqAPGSingleStage
                       := [ "StepSize",
                            "PcgsAutomorphisms",
                            "RankInitialSegmentSubgroups",
                            "SpaceEfficient",
                            "CapableDescendants",
                            "AllDescendants",
                            "Exponent",
                            "Metabelian",
                            "BasicAlgorithm",
                            "CustomiseOutput" ]
                  )
             );

#############################################################################
##
#F  AllANUPQoptions() . . . . . . . .  lists all options of the ANUPQ package
##
##  lists all the {\GAP}  options  defined  for  functions  of  the  {\ANUPQ}
##  package.
##
InstallGlobalFunction( AllANUPQoptions, function()
  return Set( Concatenation(
                  List( RecNames(ANUPQoptions), fld -> ANUPQoptions.(fld) )
                  ) );
end );

#############################################################################
##
#V  ANUPQGlobalOptions . . . . .  options that can be set globally by PqStart
##
##  A list of the options that `PqStart' can set and thereby  make  available
##  to any function  interacting  with  the  {\ANUPQ}  process  initiated  by
##  `PqStart'.
##
InstallValue( ANUPQGlobalOptions, [ "Prime", "Exponent", "Relators" ] );

#############################################################################
##
#V  ANUPQoptionChecks . . . . . . . . . . . the checks for admissible options
##
##  A record whose fields are the names of admissible ANUPQ options and whose
##  values are one-argument functions that return `true' when given  a  value
##  that is a valid value for the option, and `false' otherwise.
##
InstallValue( ANUPQoptionChecks,
              rec( Prime := x -> IsInt(x) and IsPrimeInt(x),
                   pQuotient  := IsPcGroup and IsPGroup,
                   ClassBound := IsPosInt,
                   OrderBound := IsPosInt,
                   Exponent   := IsPosInt,
                   Metabelian := IsBool,
                   GroupName  := IsString,
                   Identities := x -> IsList(x) and ForAll(x, IsFunction),
                   OutputLevel := x -> x in [0..3],
                   Relators   := x -> IsList(x) and ForAll(x, IsString),
                   StandardPresentationFile := IsString,
                   SetupFile  := IsString,
                   PqWorkspace := IsPosInt,
                   StepSize := x -> IsPosInt(x) or
                                    (IsList(x) and ForAll(x, IsPosInt)),
                   PcgsAutomorphisms := IsBool,
                   BasicAlgorithm := IsBool,
                   RankInitialSegmentSubgroups := x -> x = 0 or IsPosInt(x),
                   SpaceEfficient := IsBool,
                   CapableDescendants := IsBool,
                   AllDescendants := IsBool,
                   SubList := x -> IsPosInt(x) or
                                   (IsSet(x) and ForAll(x, IsInt)
                                    and IsPosInt(x[1])),
                   CustomiseOutput := IsRecord,
                   Bounds := x -> IsSet(x) and 2 = Length(x) and
                                  ForAll(x, IsPosInt),
                   QueueFactor := IsPosInt,
                   RedoPcp := IsBool,
                   PrintAutomorphisms := IsBool,
                   PrintPermutations := IsBool,
                   NumberOfSolubleAutomorphisms := x -> x = 0 or IsPosInt(x),
                   RelativeOrders := x -> IsList(x) and ForAll(x, IsPosInt),
                   Filename := IsString,
                   TreeDepth := IsPosInt
                   )
             );

#############################################################################
##
#V  ANUPQoptionTypes . . . . . .  the types (in words) for admissible options
##
##  A record whose fields are the names of admissible ANUPQ options and whose
##  values are valid types of the options in plain words.
##
InstallValue( ANUPQoptionTypes,
              rec( Prime := "prime integer",
                   pQuotient  := "pc p-group",
                   ClassBound := "positive integer",
                   OrderBound := "positive integer",
                   Exponent   := "positive integer",
                   Metabelian := "boolean",
                   GroupName  := "string",
                   Identities := "list of functions",
                   OutputLevel := "integer in [0..3]",
                   Relators   := "list of strings",
                   StandardPresentationFile := "string",
                   SetupFile  := "string",
                   PqWorkspace := "positive integer",
                   StepSize := "positive integer or positive integer list",
                   PcgsAutomorphisms := "boolean",
                   BasicAlgorithm := "boolean",
                   RankInitialSegmentSubgroups := "nonnegative integer",
                   SpaceEfficient := "boolean",
                   CapableDescendants := "boolean",
                   AllDescendants := "boolean",
                   SubList
                       := "pos've integer or increasing pos've integer list",
                   CustomiseOutput := "record",
                   Bounds := "pair of increasing positive integers",
                   QueueFactor := "positive integer",
                   RedoPcp := "boolean",
                   PrintAutomorphisms := "boolean",
                   PrintPermutations := "boolean",
                   NumberOfSolubleAutomorphisms := "nonnegative integer",
                   RelativeOrders := "list of positive integers",
                   Filename := "string",
                   TreeDepth := "positive integer"
                   )
             );

#############################################################################
##
#F  PQ_OTHER_OPTS_CHK( <funcname>, <interactive> ) . check opts belong to f'n
##
##  checks the `OptionsStack'  only  has  recognised  options  for  (generic)
##  function <funcname> and if not and if  `ANUPQWarnOfOtherOptions  =  true'
##  (see~"ANUPQWarnOfOtherOptions") `Info's  the  non-<funcname>  options  at
##  `InfoANUPQ' level 1.
##
##  The argument <interactive> is only relevant for those functions that have
##  both an interactive and non-interactive form, namely those with fields in
##  `PQ_FUNCTION', for which some options need to be excluded.
##
InstallGlobalFunction(PQ_OTHER_OPTS_CHK, function(funcname, interactive)
local optnames, excopts, generic, interactivestr;
  if ANUPQWarnOfOtherOptions and ValueOption("recursive") = fail and
     not IsEmpty(OptionsStack) then
    excopts := [];
    if funcname in RecNames(PQ_FUNCTION) then
      if interactive then
        excopts := ["PqWorkspace", "SetupFile"];
        interactivestr := "interactive";
      else
        interactivestr := "non-interactive";
        if funcname = "Pq" then
          excopts  := ["RedoPcp"];
        fi;
      fi;
      generic := "generic ";
    else
      generic := "";
    fi;
    optnames := Difference( RecNames( OptionsStack[ Length(OptionsStack) ] ),
                            Difference( ANUPQoptions.(funcname), excopts ) );
    if funcname = "Pq" then
      optnames := Difference( optnames, ["PqEpiOrPCover"] );
    fi;
    if not IsEmpty(optnames) then
      Info( InfoANUPQ + InfoWarning, 1,
            "ANUPQ Warning: Options: ", optnames, " ignored" );
      if IsSubset(excopts, optnames) then
        Info( InfoANUPQ + InfoWarning, 1,
              "(invalid for ", interactivestr, " call of generic function: `",
              funcname, "')." );
      else
        Info( InfoANUPQ + InfoWarning, 1,
              "(invalid for ", generic, "function: `", funcname, "')." );
      fi;
    fi;
  fi;
end);

#############################################################################
##
#F  VALUE_PQ_OPTION( <optname> ) . . . . . . . . . enhancement of ValueOption
#F  VALUE_PQ_OPTION( <optname>, <defaultval> )
#F  VALUE_PQ_OPTION( <optname>, <datarec> )
#F  VALUE_PQ_OPTION( <optname>, <defaultval>, <datarec> )
##
##  If the value <optval> of <optname> is not `fail' and it is  an  ok  value
##  for <optname> then <optval> is returned; if <optval> is not an  ok  value
##  an error is signalled. If <optval> is `fail' and <datarec> is  given  and
##  <datarec>.(<optname>) is already  bound  then  that  value  is  returned;
##  otherwise, if  <optval>  is  `fail'  and  a  default  value  <defaultval>
##  different  from  `fail'  is  supplied  then  <defaultval>  is   returned.
##  Supplying a <defaultval> of `fail' is special; it indicates  that  option
##  <optname> must have a value i.e. <optval> is not allowed to be `fail' and
##  if it is an error is signalled. If  a  <datarec>  argument  is  supplied,
##  which must be a record, then the return value, if not `fail' and a  legal
##  value, is also stored in `<datarec>.(<optname>)'.
##
##  *Note:* <defaultval> cannot be a record.
##
InstallGlobalFunction(VALUE_PQ_OPTION, function(arg)
local optname, optval, len;
  optname := arg[1];
  optval := ValueOption(optname);
  len := Length(arg);
  if optval = fail then
    if 1 = len then
      return optval;
    elif IsRecord( arg[len] ) and IsBound( arg[len].(optname) ) then
      # return the previously recorded value
      return arg[len].(optname);
    elif not IsRecord(arg[2]) then
      if arg[2] = fail then
        Error("you must supply a value for option: \"", optname, "\"\n");
      fi;
      optval := arg[2];
    fi;
  elif not ANUPQoptionChecks.(optname)(optval) then
    Error("\"", optname, "\" value must be a ",
          ANUPQoptionTypes.(optname), "\n");
  fi;
  if (optval <> fail) and (2 <= len) and IsRecord(arg[len]) then
    arg[len].(optname) := optval;
  fi;
  return optval;
end);

#############################################################################
##
#F  PQ_OPTION_CHECK(<basefn>,<datarec>) . check optns present/setable if nec.
##
##  If `<basefn> = "Pq"' (i.e. this check is  carried  out  if  the  function
##  called is `Pq', `PqEpimorphism' or  `PqPCover')  check  that  the  option
##  `Prime' has been  passed  or,  in  a  special  `PqPCover'  case,  can  be
##  determined from the `<datarec>.group'  which  must  be  present.  In  the
##  special `PqPCover' case, the options `Prime' and `ClassBound'  determined
##  are saved in <datarec>. If `Prime' is not supplied in the cases where  it
##  needs to be, an error is emitted.
##
InstallGlobalFunction(PQ_OPTION_CHECK, function(basefn, datarec)
local optname, out;
  if basefn = "Pq" then
    if datarec.calltype = "interactive" then
      if VALUE_PQ_OPTION("RedoPcp", false) then
        PQ_UNBIND(datarec, ["Prime",  "ClassBound", "Exponent", "Metabelian",
                            "pCover", "pQuotient",  "pQepi"] );
      fi;
    fi;
    if ValueOption("PqEpiOrPCover") = "pCover" and
       HasIsPGroup(datarec.group) and IsPGroup(datarec.group) then
      if VALUE_PQ_OPTION("Prime", datarec) = fail then
        if not HasPrimePGroup(datarec.group) then
          Error( "supplied group is not known to be a p-group or p unknown.\n",
                 "Option `Prime' must be supplied" );
        else
          datarec.Prime := PrimePGroup(datarec.group);
        fi;
      fi;
      if VALUE_PQ_OPTION("ClassBound", datarec) = fail and
         HasPClassPGroup(datarec.group) then
        datarec.ClassBound := PClassPGroup(datarec.group);
      fi;
    else
      VALUE_PQ_OPTION("Prime", fail, datarec);
    fi;
    VALUE_PQ_OPTION("ClassBound", 63, datarec);
  elif basefn = "StandardPresentation" then
    if VALUE_PQ_OPTION("Prime", datarec) = fail and
       VALUE_PQ_OPTION("pQuotient", datarec) = fail then
      if IsPcGroup(datarec.group) and IsPGroup(datarec.group) then
        datarec.Prime := PrimePGroup(datarec.group);
      else
        Error( "since group of process is not a pc p-group, a prime or\n",
               "p-quotient (pc group) of the group of the process ",
               "must be supplied\n" );
      fi;
    fi;
  fi;
end);

#############################################################################
##
#F  PQ_CUSTOMISE_OUTPUT(<datarec>, <subopt>, <suboptstring>, <suppstrings>)
##
##  writes the required output to the `pq' binary for the sub-option <subopt>
##  of  the  option  `CustomiseOutput',  the  value  of  that  option  having
##  previously been stored in `<datarec>.des.CustomiseOutput'; <suboptstring>
##  is part of the comment written to the `pq' binary for the sub-option  and
##  <suppstrings> is a list of such comments for the supplementary  questions
##  asked by the `pq' binary for the sub-option <subopt>.
##
InstallGlobalFunction( PQ_CUSTOMISE_OUTPUT,
function(datarec, subopt, suboptstring, suppstrings)
local optrec, isOptionSet, i;
  optrec := datarec.des.CustomiseOutput;
  if IsEmpty(suppstrings) then
    isOptionSet := IsBound( optrec.(subopt) ) and optrec.(subopt) in [1, true];
    ToPQ_BOOL(datarec, isOptionSet, suboptstring);
  elif IsBound( optrec.(subopt) ) and IsList( optrec.(subopt) ) then
    ToPQ(datarec, [ 0 ], [ "  #customise ", suboptstring ]);
    for i in [1 .. Length(suppstrings)] do
      isOptionSet := IsBound( optrec.(subopt)[i] ) and
                     optrec.(subopt)[i] in [1, true];
      ToPQ_BOOL(datarec, isOptionSet, suppstrings[i]);
    od;
  else
    ToPQ(datarec, [ 1 ], [ "  #default ", suboptstring ]);
  fi;
end);

#############################################################################
##
#F  PQ_APG_CUSTOM_OUTPUT(<datarec>, <subopt>, <suboptstring>, <suppstrings>)
##
##  writes the required output to the `pq' binary for the sub-option <subopt>
##  of the option `CustomiseOutput',  as  required  by  an  Advanced  p-Group
##  Generation Menu item, the value of that  option  having  previously  been
##  stored in `<datarec>.des.CustomiseOutput'; <suboptstring> is part of  the
##  comment written to the `pq' binary for the sub-option  and  <suppstrings>
##  is a list of such comments for the supplementary questions asked  by  the
##  `pq' binary for the sub-option <subopt>.
##
InstallGlobalFunction( PQ_APG_CUSTOM_OUTPUT,
function(datarec, subopt, suboptstring, suppstrings)
local optrec, optlist, isOptionSet, i;
  optrec := datarec.des.CustomiseOutput;
  if not( IsRecord(optrec) and IsBound( optrec.(subopt) ) and
          IsList( optrec.(subopt) ) ) then
    optlist := [];
    datarec.des.CustomiseOutput.(subopt) := optlist;
  else
    optlist := optrec.(subopt);
  fi;
  for i in [1 .. Length(suppstrings)] do
    isOptionSet := IsBound( optlist[i] ) and optlist[i] in [1, true];
    ToPQ_BOOL(datarec, isOptionSet, suppstrings[i]);
  od;
end);

#############################################################################
##
#F  SET_ANUPQ_OPTIONS( <funcname>, <fnname> )  . set options from OptionStack
##
##  When called by a function with name  <funcname>  sets  the  options  from
##  `OptionsStack'    checking    that    they    are     a     subset     of
##  `ANUPQoptions.<fnname>'. Both <funcname> and <fnname> should be strings.
##
InstallGlobalFunction( SET_ANUPQ_OPTIONS, function( funcname, fnname )
    local optrec, optnames, opt;

    # there should be options
    if IsEmpty( OptionsStack ) then
        # no options??
        optrec := rec();
        Info( InfoANUPQ, 1, funcname, " called with no options!" );
    else
        optrec := ShallowCopy( OptionsStack[ Length( OptionsStack ) ] );
        optnames := Set( REC_NAMES(optrec) );
        SubtractSet( optnames, Set( ANUPQoptions.(fnname) ) );
        Info( InfoANUPQ, 2, funcname, " called with options: ",
                            OptionsStack[ Length( OptionsStack ) ] );
        if 0 < Length(optnames) then
            # it's not an error to have unknown options,
            # function may have been called recursively and the
            # options may be intended for some other function
            Info( InfoWarning + InfoANUPQ, 2,
                  funcname, " called with unknown options: ", optnames);
        fi;
        for opt in optnames do
            Unbind( optrec.(opt) );
        od;
    fi;
    return optrec;
end );

#############################################################################
##
#F  ANUPQoptError( <funcname>, <illegal> )  . . . . . create an error message
##
##  creates an error message  for  the  function  with  name  <funcname>.  If
##  <illegal> is a string it is taken to be  the  first  line  of  the  error
##  message. Otherwise <illegal> should be alist of illegal options (strings)
##  found. The error message (string) returned also gives the list  of  valid
##  options together with the value types expected for function <funcname>.
##
InstallGlobalFunction( ANUPQoptError, function( funcname, illegal )
    local Optstring, Valstring, errmsg, optname;

    Optstring := optname -> Concatenation("\"", optname, "\"");
    Valstring := optval  -> Concatenation("<",  optval,  ">");

    if IsString(illegal) then
        errmsg := illegal;
    else # IsList(illegal)
        errmsg := Concatenation("Illegal ", funcname, " option");
        if Length(illegal) > 1 then
            Append(errmsg, "s");
        fi;
        Append(errmsg, ": ");
        Append(errmsg, JoinStringsWithSeparator( List(illegal, Optstring) ));
    fi;
    Append(errmsg, Concatenation(".\nValid ", funcname, " options:\n"));
    for optname in ANUPQoptions.(funcname) do
        Append(errmsg, "    ");
        Append(errmsg, Optstring(optname));
        if ANUPQoptionChecks.(optname) <> IsBool then
            Append(errmsg, ", ");
            Append(errmsg, Valstring( ANUPQoptionTypes.(optname) ));
        fi;
        Append(errmsg, "\n");
    od;
    return errmsg;
end );

#############################################################################
##
#F  ANUPQextractOptions( <funcname>, <args> ) . . . . . . . . extract options
##
##  extracts options from  <args>  for  function  with  name  <funcname>  and
##  returns a record suitable for use with `PushOptions'.  Abbreviations  are
##  allowed for option names so long as  each  abbreviates  a  unique  option
##  name.
##
InstallGlobalFunction( ANUPQextractOptions, function(funcname, args)
    local   Match, error, optrec, i, optname;

    # allow to give only a prefix
    Match := function( argi )
        local matches;
        if IsString(argi) then
            matches := Filtered(ANUPQoptions.(funcname),
                                optname -> 0 < Length(argi) and
                                           Length(argi) <= Length(optname) and
                                           optname{[1..Length(argi)]} = argi);
            if 1 = Length(matches) then
                return matches[1];
            fi;
            error := Concatenation( "argument: \"", argi,
                                    "\" doesn't abbreviate a unique option");
        else
            error := Concatenation( "argument: ", String(argi),
                                    " is not a (non-null) string (option name)"
                                    );
        fi;
        return fail;
    end;

    # extract options from args
    optrec := rec();
    i := 1;
    while i <= Length(args)  do
        optname := Match( args[i] );
        if optname = fail then
            Error( ANUPQoptError( funcname, error ) );
        elif ANUPQoptionChecks.(optname) = IsBool then
            optrec.(optname) := true;
            i := i + 1;
        elif i = Length(args) then
            # all remaining options are non-boolean and expect a value to
            # follow
            Error( ANUPQoptError(
                       funcname,
                       Concatenation( "Expected value for option: ", args[i] )
                       ) );
        else
            # checking values are ok is done later
            optrec.(optname) := args[i + 1];
            i := i + 2;
        fi;
    od;
    return optrec;

end );

#E  anupqopt.gi . . . . . . . . . . . . . . . . . . . . . . . . . . ends here


# 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

Declare(TPrmMulti,CodeBlock);
IsTwoPower := i >> 2 ^ Log2Int(i) = i;

get_l_power := function(t,exp)
  local p;
  if (exp=0) then
    p := L(2^t,1);
  else
    p:= L(2^t,2^(t-exp));
  fi;
  return p;
end;

# Declaration of SortBase, a 2x2 sorter.
Class(SortBase, BaseMat, rec(
   abbrevs   := [()-> []],
   new       := (self) >> SPL( WithBases(self, rec()) ).setDims(),
   dims      := self >> [ 2, 2 ],
   isReal    := True,
   sums     := self >> self,
   rChildren := self >> [],
   rSetChild := rSetChildFields(),
   toAMat := self >> Error("not supported"),
   transpose := self >> self,
));

# Declaration of SortBase, a 1x1 sorter.
Class(SortBase_w1, BaseMat, rec(
   abbrevs   := [(a)-> [a]],
   new       := (self, a) >> SPL( WithBases(self, rec(dimensions:=[1,1], a := a))),
   print := (self, i, is) >> Print(self.name, "(", self.a, ")"),
   dims      := self >> [ 1, 1 ],
   isReal    := True,
   sums     := self >> self,
   rChildren := self >> [self.a],
   rSetChild := rSetChildFields("a"),
   toAMat := self >> Error("not supported"),
   transpose := self >> self,
));

Class(SortConfigBase_w1, BaseMat, rec(
   abbrevs   := [(a)-> [a] , (b)-> [b]],
   new       := (self, a, b) >> SPL( WithBases(self, rec(dimensions:=[1,1], a := a, b:=b))),
   print := (self, i, is) >> Print(self.name, "(", self.a, ",", self.b, ")"),
   dims      := self >> [ 1, 1 ],
   isReal    := True,
   sums     := self >> self,
   rChildren := self >> [self.a, self.b],
   rSetChild := rSetChildFields("a","b"),
   toAMat := self >> Error("not supported"),
   transpose := self >> self,
));

# Declaration of SortConfigBase, a 2x2 Configurable sorter.
Class(SortConfigBase, BaseMat, rec(
   abbrevs   := [(a)-> [a]],
   new       := (self, a) >> SPL( WithBases(self, rec(dimensions:=[2,2], a := a))),
   print := (self, i, is) >> Print(self.name, "(", self.a, ")"),
   dims      := self >> [ 2, 2 ],
   isReal    := True,
   sums     := self >> self,
   rChildren := self >> [self.a],
   rSetChild := rSetChildFields("a"),
   toAMat := self >> Error("not supported"),
   transpose := self >> self,
));

# Declaration of SortLinearBase, used for linear sorter with w = 1
Class(LinearSortBase, BaseMat, rec(
   abbrevs   := [()-> []],
   new       := (self) >> SPL( WithBases(self, rec()) ).setDims(),
   dims      := self >> [ 2, 2 ],
   isReal    := True,
   sums     := self >> self,
   rChildren := self >> [],
   rSetChild := rSetChildFields(),
   toAMat := self >> Error("not supported"),
   transpose := self >> self,
));

#taken from: .../spiral/compiler/dag.gi
regassign.op_in  := self >> ConcatList(self.loc.rChildren(), ArgsExp) :: ArgsExp(self.exp);
regassign.op_out := self >> [self.loc];
regassign.op_inout := self >> [];
regassign.getNoScalar := self >> When(IsBound(self.exp.getNoScalar), self.exp.getNoScalar(), []);

# This tells Spiral how to translate 'SortBase' into code.
HDLCodegen.SortBase := (self, o, y, x, opts) >>
    chain(
	assign(nth(y,0), cond(leq(nth(x,0), nth(x,1)), nth(x,0), nth(x,1))),
	assign(nth(y,1), cond(leq(nth(x,0), nth(x,1)), nth(x,1), nth(x,0)))
    );

HDLCodegen.SortBase_w1 := (self, o, y, x, opts) >>
   let(
     t0 := TempVar(x.t.t),
     t1 := TempVar(x.t.t),
     t3 := TempVar(x.t.t),
     t5 := TempVar(x.t.t),
     t6 := TempVar(x.t.t),

     chain(
	assign(t3, imod(o.a, 2)),
	regassign(t0, cond(t3, t0, nth(x,0))),
	assign(t5, cond(leq(t0, nth(x,0)), nth(x,0), t0)),
	assign(t6, cond(leq(t0, nth(x,0)), t0, nth(x,0))),
	regassign(t1, cond(t3, t5, t1)),
	assign(nth(y,0), cond(t3, t6, t1))
     )
    );

HDLCodegen.SortConfigBase_w1 := (self, o, y, x, opts) >>
   let(
     t0 := TempVar(x.t.t),
     t1 := TempVar(x.t.t),
     t3 := TempVar(x.t.t),
     t5 := TempVar(x.t.t),
     t6 := TempVar(x.t.t),
     t7 := TempVar(x.t.t),
     t8 := TempVar(x.t.t),
     t9 := TempVar(x.t.t),
     t10 := TempVar(x.t.t),
     t2 := TempVar(x.t.t),

     chain(
	assign(t3, imod(o.a, 2)),
	regassign(t0, cond(t3, t0, nth(x,0))),
	assign(t7, eq(o.b,0)),
	assign(t8, eq(o.b,1)),

	assign(t2, leq(t0, nth(x,0))),
	assign(t9, cond(t2, t0, nth(x,0))),
	assign(t10, cond(t2, nth(x,0), t0)),
	assign(t5, cond(t7, nth(x,0) , t8, t9, t10)),
	assign(t6, cond(t7, t0 , t8, t10, t9)),

	regassign(t1, cond(t3, t5, t1)),
	assign(nth(y,0), cond(t3, t6, t1))
     )
    );

HDLCodegen.LinearSortBase := (self, o, y, x, opts) >>
   let(
     #x[1] indicates the start of a new list. When x[1] equals 1, y[0]
     #is set to the value of the register t0, and t0 is set to the value of x[0]

     t0 := TempVar(x.t.t),
     t1 := TempVar(x.t.t),
     t2 := TempVar(x.t.t),
     t3 := TempVar(x.t.t),
     t4 := TempVar(x.t.t),
     t5 := TempVar(x.t.t),
     r0 := TempVar(x.t.t),
     x0 := TempVar(x.t.t),
     x1 := TempVar(x.t.t),

     chain(

	 assign(x0, nth(x,0)),
	 assign(x1, nth(x,1)),
	 assign(t3, eq(x1, 1)),
	 regassign(r0, cond(t3, x0, t1)),
	 #regassign(r0, t0),
	 assign(t4, leq(r0, x0)),
	 assign(t1, cond(t4, x0, r0)),
	 assign(t2, cond(t4, r0, x0)),
	 assign(nth(y, 0), cond(t3, r0, t2)),
	 #regassign(t5, x1),
	 assign(nth(y, 1), x1)
     )
    );

# HDLCodegen.LinearSortBase := (self, o, y, x, opts) >>
#    let(
#      t0 := TempVar(x.t.t),
#      #x[1] indicates the start of a new list. When x[1] equals 1, y[0]
#      #is set to the value of the register t0, and t0 is set to the value of x[0]
#      chain(
# 	assign(nth(y,0),cond(eq(nth(x,1),1),t0,leq(t0, nth(x,0)),t0,nth(x,0))),
# 	regassign(t0,cond(eq(nth(x,1),1),nth(x,0),leq(t0, nth(x,0)),nth(x,0),t0)),
# 	assign(nth(y,1),nth(x,1))
#      )
#     );

HDLCodegen.SortConfigBase := (self, o, y, x, opts) >>
    let(
	t0 := TempVar(x.t.t),
	t1 := TempVar(x.t.t),
	t2 := TempVar(x.t.t),
	t3 := TempVar(x.t.t),
	chain(
	    assign(t2, nth(x,0)),
	    assign(t3, nth(x,1)),
	    assign(t0, cond(leq(t2, t3), t2, t3)),
	    assign(t1, cond(leq(t2, t3), t3, t2)),
	    assign(nth(y,0), cond(eq(o.a,0), t2, eq(o.a,1), t1, t0)),
	    assign(nth(y,1), cond(eq(o.a,0), t3, eq(o.a,1), t0, t1))
	)
    );

# Declaration of a sorter of size n
Class(Sort, TaggedNonTerminal, rec(
    abbrevs := [
    (n)       -> Checked(IsPosIntSym(n), [_unwrap(n)]),
    ],

    hashAs := self >> ObjId(self)(self.params[1]).withTags(self.getTags()),

    dims := self >> [ self.params[1], self.params[1] ],

    terminate := self >> Error("not supported"), # we could probably support this
));

# SortIJPerm(n): DirectSum(I(n/2), J(n/2))
# We do this so we can define a .permBits() function.
# This will let us easily generate large-size hardware implementations.

# (If we do not do this, the perm tool will have to manually compute the
# bit matrix representation, which will include making an n-times-n matrix.
Class(SortIJPerm, PermClass, rec(
    def := (n) -> Checked(
        IsPosIntSym(n),
        rec(size := n)),

    lambda := self >> let(
        n := self.params[1],
	fDirsum(fId(n/2), J(n/2)).lambda()
    ),

    transpose := self >> self,
    isSymmetric := self >> true,

    permBits := meth(self)
        local n, k, a, b, tmp, i;
	n := self.params[1];
	k := LogInt(n, 2);
	a := [List([1..k], i->0)];
	for i in [1 .. k-1] do
	    tmp := Concatenation([1], List([1..k-1], i->0));
	    Append(a, [tmp]);
        od;
	a := a * GF(2).one;
	b := MatSPL(I(k))*GF(2).one;

	return (a+b);
    end,
));

# Rule to breakdown Sort(n), where n is a power of two.
NewRulesFor(Sort, rec(
    Sort_Stream := rec(
        info         := "Streaming sorting network",

        applicable   := nt -> Length(nt.params) = 1 and IsTwoPower(nt.params[1]),

        children := (self, nt) >> let(

	    tag_w_tmp := nt.tags[1],
	    tag_w := tag_w_tmp.bs,
	    t := Log2Int(nt.params[1]),
	    p := Ind(2^t),
	    get_bb := w -> Cond(w=1, TTensorInd(SortBase_w1(p), p, APar, APar),TTensorI(SortBase(), 2^(t-1), APar, APar)),

	    [[ TCompose(
	           [TCompose(List([1..t-1], i ->
	               TCompose([
		           #TTensorI(SortBase(), 2^(t-1), APar, APar),
		           get_bb(tag_w),
		           TCompose(List([2..(t-i+1)], j ->
		               TCompose([
			           TTensorI(TPrm(Tensor(I(2), L(2^(j-1), 2^(j-2))) * L(2^j,2)), 2^(t-j), APar, APar),
			           get_bb(tag_w)
			           #TTensorI(SortBase(), 2^(t-1), APar, APar)
			       ])
		           )),
		        #   TTensorI(TPrm(L(2^(t-i+1), 2^(t-i)) * DirectSum(I(2^(t-i)), J(2^(t-i)))  ), 2^(i-1), APar, APar)
			   TTensorI(TPrm(L(2^(t-i+1), 2^(t-i)) * SortIJPerm(2^(t-i+1))), 2^(i-1), APar, APar)
		       ])
	            )),
		    get_bb(tag_w)]
		    #TTensorI(SortBase(), 2^(t-1), APar, APar)]
                ).withTags(nt.getTags())
            ]]
	),

        apply        := (nt, c, cnt) -> c[1],

    ),

   Sort_Stream6 := rec(
        info         := "Version of Sort_Stream (sortAlg1) removing J permutations and with configurable 2-input sorters",

        applicable   := nt -> Length(nt.params) = 1 and IsTwoPower(nt.params[1]),

        children := (self, nt) >> let(
	    t := Log2Int(nt.params[1]),
            k := Ind(2^(t-1)),
            z := (s) >> t-s,

            c1 := (s) >> logic_and(eq(bit_sel(k, z(s)), 1), neq(s, 1)),
            access_f := (s) >> cond(c1(s), 1, 2),

	    [[ TCompose(
	           [TCompose(List([1..t-1], i ->
	               TCompose([
			   TTensorInd(SortConfigBase(access_f(i)), k, APar, APar),
		           TCompose(List([2..(t-i+1)], j ->
		               TCompose([
			           TTensorI(TPrm(Tensor(I(2), L(2^(j-1), 2^(j-2))) * L(2^j,2)), 2^(t-j), APar, APar),
			           TTensorInd(SortConfigBase(access_f(i)), k, APar, APar)
			       ])
		           )),
		           TTensorI(TPrm(L(2^(t-i+1), 2^(t-i))), 2^(i-1), APar, APar)
		       ])
	            )),
		    TTensorInd(SortConfigBase(access_f(t)), k, APar, APar)]
		    #TTensorI(SortBase(), 2^(t-1), APar, APar)] #what is this line?
                ).withTags(nt.getTags())
            ]]
	),

        apply        := (nt, c, cnt) -> c[1],

    ),

   Sort_Stream_Iter := rec(
        info         := "Stream/Iter sorting network",

	depth_params := [],

        applicable   := (self, nt) >> Length(nt.params) = 1 and IsTwoPower(nt.params[1]) and
	                              Length(self.depth_params) = Log2Int(nt.params[1])-1 and
				      let(
					  t := Log2Int(nt.params[1]),
					  d := self.depth_params,
				      ForAll(List([1..t-1], i -> IsInt((t-i+1)/d[i])), j -> j)
				      ),

        children := (self, nt) >> let(
   	    t := Log2Int(nt.params[1]),

	    tag_w_tmp := nt.tags[1],
	    tag_w := tag_w_tmp.bs,
	    p := Ind(2^t),

	    get_bb := w -> Cond(w=1, TTensorInd(SortBase_w1(p), p, APar, APar),TTensorI(SortBase(), 2^(t-1), APar, APar)),

	    stage := i >> TCompose([
			     get_bb(tag_w),
			     #TTensorI(SortBase(), 2^(t-1), APar, APar),
			     TTensorI(TPrm(L(2^(t-i+1), 2^(t-i))) , 2^(i-1), APar, APar)
			  ]),

	    full_stage := i >> TCompose(List([1..self.depth_params[i]], j2 -> stage(i))),

  	    [[ TCompose(
	           [TCompose(List([1..t-1], i -> let(
		       j := Ind(t-i+1),
		       d := self.depth_params[i],
		       j1 := Ind(d),
		       j2 := Ind((t-i+1)/d),
	               TCompose([
			   Cond(d = ((t-i+1)),
                               full_stage(i),
				TCompose(List([1..d],j1 ->
				TICompose(j2, (t-i+1)/d, stage(i))
				))
			   ),
		       	   #TTensorI(TPrm(DirectSum(I(2^(t-i)), J(2^(t-i)))), 2^(i-1), APar, APar)
		       	   TTensorI(TPrm(SortIJPerm(2^(t-i+1))), 2^(i-1), APar, APar)
		       ])
	            ))),
		    get_bb(tag_w)] #last iterations left outside of TCompose
		    #TTensorI(SortBase(), 2^(t-1), APar, APar)] #last iterations left outside of TCompose
               ).withTags(nt.getTags())
            ]]
	),

        apply        := (nt, c, cnt) -> c[1],

   ),

   #old version -works-
#      Sort_Stream_Iter := rec(
#        info         := "Stream/Iter sorting network",
#
#	depth_params := [],
#
#        applicable   := (self, nt) >> Length(nt.params) = 1 and IsTwoPower(nt.params[1]) and
#	                              Length(self.depth_params) = Log2Int(nt.params[1])-1 and
#				      let(
#					  t := Log2Int(nt.params[1]),
#					  d := self.depth_params,
#				      ForAll(List([1..t-1], i -> IsInt((t-i+1)/d[i])), j -> j)
#				      ),
#
#
#        children := (self, nt) >> let(
#   	    t := Log2Int(nt.params[1]),
#
#	    stage := i >> TCompose([
#			     TTensorI(SortBase(), 2^(t-1), APar, APar),
#			     TTensorI(TPrm(L(2^(t-i+1), 2^(t-i))) , 2^(i-1), APar, APar)
#			  ]),
#
#	    full_stage := i >> TCompose(List([1..self.depth_params[i]], j2 -> stage(i))),
#
#
#  	    [[ TCompose(
#	           [TCompose(List([1..t-1], i -> let(
#		       j := Ind(t-i+1),
#		       d := self.depth_params[i],
#		       j1 := Ind((t-i+1)/d),
#	               TCompose([
#			   Cond(j1.range = 1,
#			       full_stage(i),
#			       TICompose(j1, (t-i+1)/d, full_stage(i))
#			   ),
#		       	   TTensorI(TPrm(DirectSum(I(2^(t-i)), J(2^(t-i)))), 2^(i-1), APar, APar)
#		       ])
#	            ))),
#		    TTensorI(SortBase(), 2^(t-1), APar, APar)]
#               ).withTags(nt.getTags())
#            ]]
#	),
#
#        apply        := (nt, c, cnt) -> c[1],
#
#   ),

   Sort_Stream5 := rec(
        info         := "Version of Sort_Stream_Iter (sortAlg2) removing J permutations and with configurable 2-input sorters",

	depth_params := [],

        applicable   := (self, nt) >> Length(nt.params) = 1 and IsTwoPower(nt.params[1]) and
	                              Length(self.depth_params) = Log2Int(nt.params[1])-1 and
				      let(
					  t := Log2Int(nt.params[1]),
					  d := self.depth_params,
				      ForAll(List([1..t-1], i -> IsInt((t-i+1)/d[i])), j -> j)
				      ),

        children := (self, nt) >> let(
   	    t := Log2Int(nt.params[1]),
	    k := Ind(2^(t-1)),
	    z := (s) >> t-s,

	    c1 := (s) >> logic_and(eq(bit_sel(k, z(s)), 1), neq(s, 1)),
	    access_f := (s) >> cond(c1(s), 1, 2),

	    stage := i >> TCompose([
			     TTensorInd(SortConfigBase(access_f(i)), k, APar, APar),
			     TTensorI(TPrm(L(2^(t-i+1), 2^(t-i))) , 2^(i-1), APar, APar)
			  ]),

	    full_stage := i >> TCompose(List([1..self.depth_params[i]], j2 -> stage(i))),

  	    [[ TCompose(
	           [TCompose(List([1..t-1], i -> let(
		       j := Ind(t-i+1),
		       d := self.depth_params[i],
		       j1 := Ind(d),
		       j2 := Ind((t-i+1)/d),
		       Cond(d = ((t-i+1)),
                                full_stage(i),
				TCompose(List([1..d],j1 ->
				TICompose(j2, (t-i+1)/d, stage(i))
				))
			       #unroll this change to TCompose
			)
	            ))),
		    TTensorInd(SortConfigBase(access_f(t)), k, APar, APar)] #last iterations left outside of TCompose
               ).withTags(nt.getTags())
            ]]
	),

        apply        := (nt, c, cnt) -> c[1],

   ),

  Sort_Stream3 := rec(
        info         := "",

	depth_out := 1,
        depth_in := [],

        applicable   := (self, nt) >> Length(nt.params) = 1 and IsTwoPower(nt.params[1]) and
			     	      Length(self.depth_in) = Log2Int(nt.params[1])-1  and
				      let(
                                          t := Log2Int(nt.params[1]), d_out := self.depth_out,
                                          d_in := self.depth_in, IsInt(t/d_out) and ((d_out=1) or ((d_out=t) and
                                      	  (ForAll(List([1..t-1], i -> IsInt((t-i+1)/d_in[i])), j -> j))))
                                      ),

        children := (self, nt) >> let(
	    d_out := self.depth_out,
	    d_in := self.depth_in,
            t := Log2Int(nt.params[1]),
	    k := Ind(2^(t-1)),
	    iter_d1 := Ind(t),

            z := (lp, jp) >> (t-1)-(lp+jp),
            c2 := (lp, jp) >> logic_and(eq(bit_sel(k, z(lp, jp)), 1), neq(lp, 0)),
            access_f := (lp, jp) >> cond(c2(lp, jp), 1, 2),
            stage := (lp, jp) >> TCompose([
                       TTensorInd(SortConfigBase(access_f(lp, jp)), k, APar, APar),
                       TPrm(L(2^t, 2^(t-1)))
                   ]),
	    full_stage := i >> TCompose(List([0..d_in[i+1]-1], j -> stage(i,j))),
	    [[ Cond(d_out=1,
			TICompose(iter_d1,t,let(
				iter_d2 := Ind(t-iter_d1),
	    			p_list := List([0..t-1], i-> get_l_power(t,i)),
				TCompose([
					TICompose(iter_d2,t-iter_d1, stage(iter_d1,iter_d2)),
		  	 		TPrmMulti(p_list,iter_d1)
				])
			)).withTags(nt.getTags()),
		    d_out=t,
			TCompose([
				TCompose(List([0..t-2], iter_d3 -> let(
				  d_stage := d_in[iter_d3+1],
            			  iter_d5 := Ind((t-iter_d3)/d_stage),
				  TCompose([
					Cond(d_stage=(t-iter_d3),
						full_stage(iter_d3),
						TCompose(List([0..d_stage-1],iter_d4 ->
						  TICompose(iter_d5,(t-iter_d3)/d_stage,stage(iter_d3,iter_d4*iter_d5))))
					),
					get_l_power(t,iter_d3)
				  ])
				))),
				stage(t-1,0),
				get_l_power(t,t-1)
			]).withTags(nt.getTags())
	    )]]

	),
        apply        := (nt, c, cnt) -> c[1],
   ),

   Sort_Stream4 := rec(
        info         := "",

	depth := 1,

        applicable   := (self, nt) >> Length(nt.params) = 1 and IsTwoPower(nt.params[1]) and
	                              let (d := self.depth, t := Log2Int(nt.params[1]), IsInt(t*t/self.depth) and
				      (IsInt(d/t) or IsInt(t/d))),

        children := (self, nt) >> let(
   	    t := Log2Int(nt.params[1]),
	    d := self.depth,
	    d1 := cond(leq(d,t), d, t).ev(),
	    d2 := cond(leq(d,t), 1, d/t).ev(),

	    k := Ind(2^(t-1)),
	    j := Ind(t),
	    l := Ind(t),
	    n := Ind(t/d2),
	    v := Ind(t/d1),

            d_tmp := cond(leq(d,t), t/d, 0).ev(),
            d2_tmp := cond(leq(t,d), d/t, 0).ev(),
	    s_tmp := ((t*t)/d),
            m2 := Ind(d_tmp),
	    s2 := Ind(s_tmp),
	    v2 := Ind(d2_tmp),
	    n2 := Ind(d),
	    n3 := Ind(d),
	    l2 := Ind(t),

	    tag_w_tmp := nt.tags[1],
	    tag_w := tag_w_tmp.bs,
	    p := Ind(2^t),

	    c1 := (lp, jp) >> lt((t-1), (lp+jp)),
	    z := (lp, jp) >> (t-1)-(lp+jp),
	    z_w1 := (lp, jp) >> (t-1)-(lp+jp)+1,
	    c2 := (lp, jp) >> logic_and(eq(bit_sel(k, z(lp, jp)), 1), neq(lp, 0)),
	    c2_w1 := (lp, jp) >> logic_and(eq(bit_sel(p, z_w1(lp, jp)), 1), neq(lp, 0)),

	    access_f := (lp, jp) >> cond(c1(lp, jp), 0, c2(lp, jp), 1, 2),
	    access_f_w1 := (lp, jp) >> cond(c1(lp, jp), 0, c2_w1(lp, jp), 1, 2),

	    get_bb := (lp,jp) -> Cond(tag_w=1, TTensorInd(SortConfigBase_w1(p,access_f_w1(lp,jp)), p, APar, APar),TTensorInd(SortConfigBase(access_f(lp, jp)), k, APar, APar)),

	    stage := (lp, jp) >> TCompose([
		       get_bb(lp, jp),
		       #TTensorInd(SortConfigBase(access_f(lp, jp)), k, APar, APar),
		       TPrm(L(2^t, 2^(t-1)))
	           ]),

	    full_stage := np_1 >> TCompose(List([0..t-1], m_1 -> TCompose(List([0..t-1], j_1 -> stage(m_1, j_1))))),
	    full_stage1 := np >> TCompose(List([0..d2-1], m -> TCompose(List([0..t-1], j -> stage(d2*np+m, j))))),
	    full_stage2 := vp >> TCompose(List([0..d1-1], s -> stage(l, vp+s))),
            full_stage1b := np2 >> TICompose(m2,d_tmp, TICompose(j,t, stage((t/d)*np2+m2, j))),

             # Old: problem is that it's assuming the l2 above, which is an unassigned iterator.
	     # There is also a problem with the vp2+s2 parameter: you need to multiply vp2 by the number of iterations.
             # full_stage2b := (vp2) >> TICompose(s2,s_tmp, stage(l2, vp2+s2)),
            full_stage2b := (vp2, l3) >> TICompose(s2,s_tmp, stage(l3, vp2*s_tmp+s2)),

  	    [[ Cond(d=t*t, full_stage(0).withTags(nt.getTags()),
		    d<t, TCompose(List([0..d-1], n2 -> full_stage1b(n2))).withTags(nt.getTags()),
                    d=t, TCompose(List([0..d-1], n3 -> full_stage1b(n3))).withTags(nt.getTags()),
                    d>t, TCompose(List([0..t-1], l3 -> TCompose(List([0..d2_tmp-1], v2 -> full_stage2b(v2, l3))))).withTags(nt.getTags())) #the problem seems to be the outer most TCompose works if it was TICompose
		    #d>t, TICompose(n, t/d2, full_stage1(n)).withTags(nt.getTags())) #old one but will leave it as new one does not work yet
	        ]]
	),

#	Old version -works-
#	   Sort_Stream4 := rec(
#        info         := "",
#
#	depth := 1,
#
#        applicable   := (self, nt) >> Length(nt.params) = 1 and IsTwoPower(nt.params[1]) and
#	                              let (d := self.depth, t := Log2Int(nt.params[1]), IsInt(t*t/self.depth) and
#				      (IsInt(d/t) or IsInt(t/d))),
#
#
#        children := (self, nt) >> let(
#   	    t := Log2Int(nt.params[1]),
#	    d := self.depth,
#	    d1 := cond(leq(d,t), d, t).ev(),
#	    d2 := cond(leq(d,t), 1, d/t).ev(),
#
#	    k := Ind(2^(t-1)),
#	    j := Ind(t),
#	    l := Ind(t),
#	    n := Ind(t/d2),
#	    v := Ind(t/d1),
#
#
#	    c1 := (lp, jp) >> lt((t-1), (lp+jp)),
#	    z := (lp, jp) >> (t-1)-(lp+jp),
#	    c2 := (lp, jp) >> logic_and(eq(bit_sel(k, z(lp, jp)), 1), neq(lp, 0)),
#
#	    access_f := (lp, jp) >> cond(c1(lp, jp), 0, c2(lp, jp), 1, 2),
#
#	    stage := (lp, jp) >> TCompose([
#		       TTensorInd(SortConfigBase(access_f(lp, jp)), k, APar, APar),
#		       TPrm(L(2^t, 2^(t-1)))
#	           ]),
#
#	    full_stage1 := np >> TCompose(List([0..d2-1], m -> TCompose(List([0..t-1], j -> stage(d2*np+m, j))))),
#	    full_stage2 := vp >> TCompose(List([0..d1-1], s -> stage(l, vp+s))),
#
#  	    [[ Cond(d=t*t, full_stage1(0).withTags(nt.getTags()),
#		    d>t, TICompose(n, t/d2, full_stage1(n)).withTags(nt.getTags()),
#		    d<t, TICompose(l, t, TICompose(v, t/d1, full_stage2(d1*v))).withTags(nt.getTags()),
#		    d=t, TICompose(l, t, full_stage2(0)).withTags(nt.getTags()))
#	        ]]
#	),
#
#        apply        := (nt, c, cnt) -> c[1],
#   ),

	# this works, but only for d=1.
#         children := (self, nt) >> let(
#    	    t := Log2Int(nt.params[1]),
# 	    k := Ind(2^(t-1)),
# 	    l := Ind(t),
# 	    j := Ind(t),

# 	    c1 := lt((t-1), (l+j)),
# 	    z := (t-1)-(l+j),
# 	    c2 := logic_and(eq(bit_sel(k, z), 1), neq(l, 0)),

# 	    access_f := cond(c1, 0, c2, 1, 2),

# 	    d1 := cond(leq(d,t), d, t),
# 	    d2 := cond(leq(d,t), 1, d/t),

# 	    c1 := lt((t-1), (l+j)),
# 	    z := (t-1)-(l+j),
# 	    c2 := logic_and(eq(bit_sel(k, z), 1), neq(l, 0)),

# 	    access_f := cond(c1, 0, c2, 1, 2),

# 	    d1 := cond(leq(d,t), d, t),
# 	    d2 := cond(leq(d,t), 1, d/t),

#   	    [[ TICompose(l, t,
# 		   TICompose(j, t, TCompose([
# 		       TTensorInd(SortConfigBase(access_f), k, APar, APar),
# 		       TPrm(TL(2^t, 2^(t-1), 1, 1))
# 	           ]))).withTags(nt.getTags())
#             ]]
# 	),

        apply        := (nt, c, cnt) -> c[1],

   ),

  Linear_Sort := rec(
        info         := "",

        applicable   := (self, nt) >> IsTwoPower(nt.params[1]),

        children := (self, nt) >> let(
            t := Log2Int(nt.params[1]),
	    [[ TCompose(List([1..2^t], i ->
		TTensorI(LinearSortBase(), nt.params[1], APar, APar)
			)).withTags(nt.getTags()),
	    ]]

	),
        apply        := (nt, c, cnt) -> c[1],
   ),

));

#---------------------------------------------------------------------
#---------------------------------------------------------------------
#---------------------------------------------------------------------
# Old stuff, not used now

#     Sort_Stream_old := rec(
#         info         := "Streaming sorting network",

#         switch       := true,
#         applicable   := nt ->
#             Length(nt.params) = 1 and IsTwoPower(nt.params[1]),

#         children := (self, nt) >> let(
# 	    t := Log2Int(nt.params[1]),
# 	    [[ TCompose(List([1..t], i ->
# 	 	TCompose([
# 		    TTensorI(BitonicSort(2^(t-i+1)), 2^(i-1), APar, APar),
# 		    TTensorI(TPrm(DirectSum(I(2^(t-i)), J(2^(t-i)))), 2^(i-1), APar, APar)
# 		])
#             )).withTags(nt.getTags())]]
#         ),

#         apply        := (nt, c, cnt) -> c[1],

#     ),

# Declaration of a bitonic sorter of size n
# Class(BitonicSort, TaggedNonTerminal, rec(
#     abbrevs := [
#     (n)       -> Checked(IsPosIntSym(n), [_unwrap(n)]),
#     ],

#     hashAs := self >> ObjId(self)(self.params[1]).withTags(self.getTags()),

#     dims := self >> [ self.params[1], self.params[1] ],

#     terminate := self >> Error("not supported"), # we could probably support this
# ));

# Old stuff, not used now.
# # Rules to break-down BitonicSort()
# NewRulesFor(BitonicSort, rec(
#     BitonicSort_Stream := rec(
#         info         := "Streaming Bitonic sorting network",

#         switch       := true,
#         applicable   := nt ->
#             Length(nt.params) = 1 and IsTwoPower(nt.params[1]) and nt.params[1] > 2,

#         children := (self, nt) >> let(
# 	    k := Log2Int(nt.params[1]),
# 	    [[ TCompose([ TTensorI(SortBase(), 2^(k-1), APar, APar),
# 		  TCompose(List([2..k], j -> TCompose([
# 		      TTensorI(TPrm(Tensor(I(2), L(2^(j-1), 2^(j-2))) * L(2^j, 2)), 2^(k-j), APar, APar),
# 		      TTensorI(SortBase(), 2^(k-1), APar, APar)
# 		  ]))),
# 		  TPrm(L(2^k, 2^(k-1)))
# 	       ]).withTags(nt.getTags())
# 	    ]]
#         ),

#         apply        := (nt, c, cnt) -> c[1],
#     ),

#     BitonicSort_Base := rec(
# 	info := "Bitonic soring network base rule",
# 	applicable := nt -> nt.params[1] = 2,
# 	chldren := (self, nt) >> [[ ]],
# 	apply := (nt, c, cnt) -> SortBase()
#     )

# ));


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

Import(search);

#   These functions merge multiple hashtables into a single one
AllHashEntries := function(table)
    local e, i, res;
    res := [];
    for e in Filtered(table.entries, IsList) do
        for i in e do
            if i.data <> [] then Add(res, i); fi;
        od;
    od;
    return res;
end;

AddMergedHashTables := function(target, newhashs)
    local e, v, h, i;
    for e in Flat(List(newhashs, i->AllHashEntries(i))) do
        v := [];
        for h in newhashs do
            Add(v, HashLookup(h, e.key));
        od;
        v := Filtered(v, i -> i <> [] and i <> false);
        if v <> [] then
            Sort(v, (j,k) -> j[1].measured < k[1].measured);
            HashAdd(target, e.key, v[1]);
        fi;
    od;
end;

#   Build all base cases for a given SIMD ISA.
#   Currently only L(2v, 2), L(2v, v), L(v^2, v)
#
SIMD_ISA_DB.buildBases := meth(self, isa)
    local rebind, t, v, tags, cxtags, t1, t2, rt, brules, common, rset1, rset2, tab1, tab2, e, h;

    if self.verbose then Print("Building bases for ", isa, "\n"); fi;
    self.hash_rebuilt := true;
    v := isa.v;
    tags := isa.getTags();
    cxtags := isa.getTagsCx();

    # == do TL base cases======================================
    # TL usually is not measured in DP
    rebind := IsBound(TL.doNotMeasure) and TL.doNotMeasure;
    if rebind then
        TL.doNotMeasure := false;
    fi;

    brules := paradigms.tSPL_Globals.getDPOpts().breakdownRules;
    common := [ SIMD_ISA_Bases1, SIMD_ISA_Bases2, IxLxI_kmn_n, IxLxI_kmn_km, paradigms.vector.breakdown.IxLxI_vtensor ];

    rset1 := CopyFields(brules, rec(TL := Concat(common, [IxLxI_IxLxI_up])));
    rset2 := CopyFields(brules, rec(TL := Concat(common, [IxLxI_IxLxI_down])));

    tab1 := HashTableDP();
    for t in SIMD_ISA_DB.getBases(isa) do
#    [ TL(2*v,v,1,1).withTags(tags), TL(2*v,2,1,1).withTags(tags), TL(v*v,v,1,1).withTags(tags), TL(v*v/4,v/2,1,2).withTags(tags) ]
        t1 := DP(t, rec(measureFunction := VCost, verbosity := 0, hashTable := tab1, globalUnrolling := true),
                CopyFields(isa.splopts, rec(breakdownRules := rset1, dataType := "no default", baseHashes := [], globalUnrolling := 10000)));
    od;

    tab2 := HashTableDP();
    for t in SIMD_ISA_DB.getBases(isa) do
#    [ TL(2*v,v,1,1).withTags(tags), TL(2*v,2,1,1).withTags(tags), TL(v*v,v,1,1).withTags(tags), TL(v*v/4,v/2,1,2).withTags(tags) ]
        t2 := DP(t, rec(measureFunction := VCost, verbosity := 0, hashTable := tab2, globalUnrolling := false),
                CopyFields(isa.splopts, rec(breakdownRules := rset2, dataType := "no default", baseHashes := [], globalUnrolling := 10000)));
    od;
    AddMergedHashTables(self.hash, [tab1, tab2]);
    # restore TL
    if rebind then
        TL.doNotMeasure:=true;
    fi;

    # == do other base cases===================================
    # do other base cases
    #
    # none here yet :(
    #

    # == final check - did all work? ==========================
    for e in self.getBases(isa) do
        h := HashLookup(self.hash, e);
        if  h=false or h=[] then Print(e, " could not be built\n"); fi;
    od;
end;