Implement blockget removal for known (constant) blocks

compiler/src/l3/CPSOptimizer.scala

@@ -20,7 +20,8 @@ abstract class CPSOptimizer[T <: CPSTreeModule { type Name = Symbol }]
     cSubst: Subst[Name] = emptySubst,
     eInvEnv: Map[(ValuePrimitive, Seq[Atom]), Atom] = Map.empty,
     cEnv: Map[Name, Cnt] = Map.empty,
-    fEnv: Map[Name, Fun] = Map.empty) {
+    fEnv: Map[Name, Fun] = Map.empty,
+    bEnv: Map[(Atom, Atom), Atom] = Map.empty) {
 
     def eEnv: Map[Atom, (ValuePrimitive, Seq[Atom])] =
       eInvEnv.map(_.swap)
@@ -50,6 +51,9 @@ abstract class CPSOptimizer[T <: CPSTreeModule { type Name = Symbol }]
       copy(cEnv = cEnv ++ (cnts.map(_.name) zip cnts))
     def withFuns(funs: Seq[Fun]): State =
       copy(fEnv = fEnv ++ (funs.map(_.name) zip funs))
+
+    def withBlockVal(block: Atom, index: Atom, value: Atom): State =
+      copy(bEnv = bEnv + ((block, index) -> value))
   }
 
   // Shrinking optimizations
@@ -142,6 +146,26 @@ abstract class CPSOptimizer[T <: CPSTreeModule { type Name = Symbol }]
           // Right Absorbing
           case Seq(a1, AtomL(l2)) if rightAbsorbing((prim, l2)) =>
             shrink(body, s.withASubst(name, l2))
+
+          case Seq(b, idx) if prim == blockGet && s.bEnv.isDefinedAt((b, idx)) =>
+            shrink(body, s.withASubst(name, s.bEnv((b, idx))))
+          case Seq(b, idx, v) if prim == blockSet && s.eEnv.isDefinedAt(b) =>
+            val maybeBlockAlloc = s.eEnv(b)._1
+            val isKnownConstant = blockAllocTag
+              .lift(maybeBlockAlloc)
+              .filter(constantTags(_))
+              .isDefined
+
+            if (isKnownConstant) {
+              val newBody = shrink(body, newState.withBlockVal(b, idx, v))
+              LetP(name, prim, replacedArgs, newBody)
+            } else {
+              noOp
+            }
+
+          // Single argument primitives
+          // currently only checks for block length and
+          // block tag
           case Seq(a1) if s.eEnv.isDefinedAt(a1) =>
             val (maybeBlockAlloc, maybeLengthAtom) = s.eEnv(a1)
             blockAllocTag.lift(maybeBlockAlloc) match {
@@ -381,6 +405,9 @@ abstract class CPSOptimizer[T <: CPSTreeModule { type Name = Symbol }]
   protected val blockAllocTag: PartialFunction[ValuePrimitive, Literal]
   protected val blockTag: ValuePrimitive
   protected val blockLength: ValuePrimitive
+  protected val blockGet: ValuePrimitive
+  protected val blockSet: ValuePrimitive
+  protected val constantTags: Set[Literal]
 
   protected val identity: ValuePrimitive
 
@@ -423,6 +450,9 @@ object CPSOptimizerHigh extends CPSOptimizer(SymbolicCPSTreeModule)
   }
   protected val blockTag: ValuePrimitive = BlockTag
   protected val blockLength: ValuePrimitive = BlockLength
+  protected val blockGet: ValuePrimitive = BlockGet
+  protected val blockSet: ValuePrimitive = BlockSet
+  protected val constantTags: Set[Literal] = Set(intToLit(200), intToLit(202))
 
   protected val identity: ValuePrimitive = Id
 
@@ -520,6 +550,9 @@ object CPSOptimizerLow extends CPSOptimizer(SymbolicCPSTreeModuleLow)
   }
   protected val blockTag: ValuePrimitive = BlockTag
   protected val blockLength: ValuePrimitive = BlockLength
+  protected val blockGet: ValuePrimitive = BlockGet
+  protected val blockSet: ValuePrimitive = BlockSet
+  protected val constantTags = Set(200, 202)
 
   protected val identity: ValuePrimitive = Id