Shake bug fix; pattern reduction in ReduceLocal

5 files changed, 116 insertions, 16 deletions

diff --git a/src/compiler.sig b/src/compiler.sig
index 28a5a5d5..1a41eaea 100644
--- a/src/compiler.sig
+++ b/src/compiler.sig
@@ -115,14 +115,17 @@ signature COMPILER = sig
     val toRpcify : (string, Core.file) transform
     val toCore_untangle2 : (string, Core.file) transform
     val toShake2 : (string, Core.file) transform
+    val toEspecialize1 : (string, Core.file) transform 
+    val toCore_untangle3 : (string, Core.file) transform
+    val toShake3 : (string, Core.file) transform
     val toTag : (string, Core.file) transform
     val toReduce : (string, Core.file) transform
     val toUnpoly : (string, Core.file) transform 
     val toSpecialize : (string, Core.file) transform 
-    val toShake3 : (string, Core.file) transform
-    val toEspecialize : (string, Core.file) transform 
-    val toReduce2 : (string, Core.file) transform
     val toShake4 : (string, Core.file) transform
+    val toEspecialize2 : (string, Core.file) transform 
+    val toReduce2 : (string, Core.file) transform
+    val toShake5 : (string, Core.file) transform
     val toMarshalcheck : (string, Core.file) transform
     val toEffectize : (string, Core.file) transform
     val toMonoize : (string, Mono.file) transform
diff --git a/src/compiler.sml b/src/compiler.sml
index c8059c6e..256162ce 100644
--- a/src/compiler.sml
+++ b/src/compiler.sml
@@ -753,13 +753,16 @@ val toRpcify = transform rpcify "rpcify" o toShake1
 
 val toCore_untangle2 = transform core_untangle "core_untangle2" o toRpcify
 val toShake2 = transform shake "shake2" o toCore_untangle2
+val toEspecialize1 = transform especialize "especialize1" o toShake2
+val toCore_untangle3 = transform core_untangle "core_untangle3" o toEspecialize1
+val toShake3 = transform shake "shake3" o toCore_untangle3
 
 val tag = {
     func = Tag.tag,
     print = CorePrint.p_file CoreEnv.empty
 }
 
-val toTag = transform tag "tag" o toCore_untangle2
+val toTag = transform tag "tag" o toShake3
 
 val reduce = {
     func = Reduce.reduce,
@@ -782,20 +785,20 @@ val specialize = {
 
 val toSpecialize = transform specialize "specialize" o toUnpoly
 
-val toShake3 = transform shake "shake3" o toSpecialize
+val toShake4 = transform shake "shake4" o toSpecialize
 
-val toEspecialize = transform especialize "especialize" o toShake3
+val toEspecialize2 = transform especialize "especialize2" o toShake4
 
-val toReduce2 = transform reduce "reduce2" o toEspecialize
+val toReduce2 = transform reduce "reduce2" o toEspecialize2
 
-val toShake4 = transform shake "shake4" o toReduce2
+val toShake5 = transform shake "shake5" o toReduce2
 
 val marshalcheck = {
     func = (fn file => (MarshalCheck.check file; file)),
     print = CorePrint.p_file CoreEnv.empty
 }
 
-val toMarshalcheck = transform marshalcheck "marshalcheck" o toShake4
+val toMarshalcheck = transform marshalcheck "marshalcheck" o toShake5
 
 val effectize = {
     func = Effective.effectize,
diff --git a/src/reduce_local.sml b/src/reduce_local.sml
index a9f28617..4ddddfbf 100644
--- a/src/reduce_local.sml
+++ b/src/reduce_local.sml
@@ -33,6 +33,12 @@ open Core
 
 structure IM = IntBinaryMap
 
+fun multiLiftExpInExp n e =
+    if n = 0 then
+        e
+    else
+        multiLiftExpInExp (n - 1) (CoreEnv.liftExpInExp 0 e)
+
 datatype env_item =
          Unknown
        | Known of exp
@@ -44,6 +50,76 @@ type env = env_item list
 val deKnown = List.filter (fn Known _ => false
                             | _ => true)
 
+datatype result = Yes of env | No | Maybe
+
+fun match (env, p : pat, e : exp) =
+    let
+        val baseline = length env
+
+        fun match (env, p, e) =
+            case (#1 p, #1 e) of
+                (PWild, _) => Yes env
+              | (PVar (x, t), _) => Yes (Known (multiLiftExpInExp (length env - baseline) e) :: env)
+
+              | (PPrim p, EPrim p') =>
+                if Prim.equal (p, p') then
+                    Yes env
+                else
+                    No
+
+              | (PCon (_, PConVar n1, _, NONE), ECon (_, PConVar n2, _, NONE)) =>
+                if n1 = n2 then
+                    Yes env
+                else
+                    No
+
+              | (PCon (_, PConVar n1, _, SOME p), ECon (_, PConVar n2, _, SOME e)) =>
+                if n1 = n2 then
+                    match (env, p, e)
+                else
+                    No
+
+              | (PCon (_, PConFfi {mod = m1, con = con1, ...}, _, NONE),
+                 ECon (_, PConFfi {mod = m2, con = con2, ...}, _, NONE)) =>
+                if m1 = m2 andalso con1 = con2 then
+                    Yes env
+                else
+                    No
+
+              | (PCon (_, PConFfi {mod = m1, con = con1, ...}, _, SOME ep),
+                 ECon (_, PConFfi {mod = m2, con = con2, ...}, _, SOME e)) =>
+                if m1 = m2 andalso con1 = con2 then
+                    match (env, p, e)
+                else
+                    No
+
+              | (PRecord xps, ERecord xes) =>
+                if List.exists (fn ((CName _, _), _, _) => false
+                                 | _ => true) xes then
+                    Maybe
+                else
+                    let
+                        fun consider (xps, env) =
+                            case xps of
+                                [] => Yes env
+                              | (x, p, _) :: rest =>
+                                case List.find (fn ((CName x', _), _, _) => x' = x
+                                                 | _ => false) xes of
+                                    NONE => No
+                                  | SOME (_, e, _) =>
+                                    case match (env, p, e) of
+                                        No => No
+                                      | Maybe => Maybe
+                                      | Yes env => consider (rest, env)
+                    in
+                        consider (xps, env)
+                    end
+
+              | _ => Maybe
+    in
+        match (env, p, e)
+    end
+
 fun exp env (all as (e, loc)) =
     case e of
         EPrim _ => all
@@ -127,11 +203,24 @@ fun exp env (all as (e, loc)) =
                   | PCon (_, _, _, NONE) => 0
                   | PCon (_, _, _, SOME p) => patBinds p
                   | PRecord xpts => foldl (fn ((_, p, _), n) => n + patBinds p) 0 xpts
+
+            fun push () =
+                (ECase (exp env e,
+                        map (fn (p, e) => (p,
+                                           exp (List.tabulate (patBinds p,
+                                                            fn _ => Unknown) @ env) e))
+                            pes, others), loc)
+
+            fun search pes =
+                case pes of
+                    [] => push ()
+                  | (p, body) :: pes =>
+                    case match (env, p, e) of
+                        No => search pes
+                      | Maybe => push ()
+                      | Yes env' => exp env' body
         in
-            (ECase (exp env e,
-                    map (fn (p, e) => (p,
-                                       exp (List.tabulate (patBinds p, fn _ => Unknown) @ env) e))
-                        pes, others), loc)
+            search pes
         end
 
       | EWrite e => (EWrite (exp env e), loc)
diff --git a/src/shake.sml b/src/shake.sml
index 291f2fb0..dde131fc 100644
--- a/src/shake.sml
+++ b/src/shake.sml
@@ -67,7 +67,7 @@ fun shake file =
 
         val (usedE, usedC) =
             List.foldl
-                (fn ((DExport (_, n), _), (usedE, usedC)) => (IS.add (usedE, n), usedE)
+                (fn ((DExport (_, n), _), (usedE, usedC)) => (IS.add (usedE, n), usedC)
                   | ((DTable (_, _, c, _, pe, pc, ce, cc), _), (usedE, usedC)) =>
                     let
                         val usedC = usedVarsC usedC c
@@ -170,7 +170,7 @@ fun shake file =
 
         val s = IS.foldl (fn (n, s) =>
                              case IM.find (cdef, n) of
-                                 NONE => raise Fail "Shake: Couldn't find 'con'"
+                                 NONE => raise Fail ("Shake: Couldn't find 'con' " ^ Int.toString n)
                                | SOME cs => foldl (fn (c, s) => shakeCon s c) s cs) s usedC
     in
         List.filter (fn (DCon (_, n, _, _), _) => IS.member (#con s, n)
diff --git a/src/tag.sml b/src/tag.sml
index 582a3b8e..b4574b79 100644
--- a/src/tag.sml
+++ b/src/tag.sml
@@ -74,6 +74,8 @@ fun exp env (e, s) =
                                            let
                                                fun tagIt (ek, newAttr) =
                                                    let
+                                                       val eOrig = e
+
                                                        fun unravel (e, _) =
                                                            case e of
                                                                ENamed n => (n, [])
@@ -83,7 +85,10 @@ fun exp env (e, s) =
                                                                in
                                                                    (n, es @ [e2])
                                                                end
-                                                             | _ => (ErrorMsg.errorAt loc "Invalid link expression";
+                                                             | _ => (ErrorMsg.errorAt loc ("Invalid " ^ newAttr
+                                                                                           ^ " expression");
+                                                                     Print.epreface ("Expression",
+                                                                                     CorePrint.p_exp CoreEnv.empty eOrig);
                                                                      (0, []))
 
                                                        val (f, args) = unravel e