Especialize

10 files changed, 267 insertions, 31 deletions

diff --git a/src/compiler.sig b/src/compiler.sig
index 0c95934a..e26ec13c 100644
--- a/src/compiler.sig
+++ b/src/compiler.sig
@@ -61,6 +61,7 @@ signature COMPILER = sig
     val termination : (Elab.file, Elab.file) phase
     val explify : (Elab.file, Expl.file) phase
     val corify : (Expl.file, Core.file) phase
+    val especialize : (Core.file, Core.file) phase
     val shake : (Core.file, Core.file) phase
     val tag : (Core.file, Core.file) phase
     val reduce : (Core.file, Core.file) phase
@@ -82,6 +83,7 @@ signature COMPILER = sig
     val toTermination : (string, Elab.file) transform
     val toExplify : (string, Expl.file) transform
     val toCorify : (string, Core.file) transform
+    val toEspecialize : (string, Core.file) transform 
     val toShake1 : (string, Core.file) transform
     val toTag : (string, Core.file) transform
     val toReduce : (string, Core.file) transform
diff --git a/src/compiler.sml b/src/compiler.sml
index 1f88705c..4f1bce11 100644
--- a/src/compiler.sml
+++ b/src/compiler.sml
@@ -404,12 +404,19 @@ val corify = {
 
 val toCorify = transform corify "corify" o toExplify
 
+val especialize = {
+    func = ESpecialize.specialize,
+    print = CorePrint.p_file CoreEnv.empty
+}
+
+val toEspecialize = transform especialize "especialize" o toCorify
+
 val shake = {
     func = Shake.shake,
     print = CorePrint.p_file CoreEnv.empty
 }
 
-val toShake1 = transform shake "shake1" o toCorify
+val toShake1 = transform shake "shake1" o toEspecialize
 
 val tag = {
     func = Tag.tag,
diff --git a/src/core_env.sig b/src/core_env.sig
index cdbf5946..98e345cc 100644
--- a/src/core_env.sig
+++ b/src/core_env.sig
@@ -33,6 +33,9 @@ signature CORE_ENV = sig
     val liftConInExp : int -> Core.exp -> Core.exp
     val subConInExp : (int * Core.con) -> Core.exp -> Core.exp
 
+    val liftExpInExp : int -> Core.exp -> Core.exp
+    val subExpInExp : (int * Core.exp) -> Core.exp -> Core.exp
+
     type env
 
     val empty : env
diff --git a/src/core_env.sml b/src/core_env.sml
index b399f62f..0faf5aab 100644
--- a/src/core_env.sml
+++ b/src/core_env.sml
@@ -93,6 +93,35 @@ val subConInExp =
                 bind = fn ((xn, rep), U.Exp.RelC _) => (xn+1, liftConInCon 0 rep)
                         | (ctx, _) => ctx}
 
+val liftExpInExp =
+    U.Exp.mapB {kind = fn k => k,
+                con = fn _ => fn c => c,
+                exp = fn bound => fn e =>
+                                     case e of
+                                         ERel xn =>
+                                         if xn < bound then
+                                             e
+                                         else
+                                             ERel (xn + 1)
+                                       | _ => e,
+                bind = fn (bound, U.Exp.RelE _) => bound + 1
+                        | (bound, _) => bound}
+
+val subExpInExp =
+    U.Exp.mapB {kind = fn k => k,
+                con = fn _ => fn c => c,
+                exp = fn (xn, rep) => fn e =>
+                                  case e of
+                                      ERel xn' =>
+                                      (case Int.compare (xn', xn) of
+                                           EQUAL => #1 rep
+                                         | GREATER=> ERel (xn' - 1)
+                                         | LESS => e)
+                                    | _ => e,
+                bind = fn ((xn, rep), U.Exp.RelE _) => (xn+1, liftExpInExp 0 rep)
+                        | ((xn, rep), U.Exp.RelC _) => (xn, liftConInExp 0 rep)
+                        | (ctx, _) => ctx}
+
 (* Back to environments *)
 
 exception UnboundRel of int
diff --git a/src/core_util.sig b/src/core_util.sig
index 43750698..2ae75305 100644
--- a/src/core_util.sig
+++ b/src/core_util.sig
@@ -107,6 +107,11 @@ structure Exp : sig
     val exists : {kind : Core.kind' -> bool,
                   con : Core.con' -> bool,
                   exp : Core.exp' -> bool} -> Core.exp -> bool
+
+    val foldMap : {kind : Core.kind' * 'state -> Core.kind' * 'state,
+                   con : Core.con' * 'state -> Core.con' * 'state,
+                   exp : Core.exp' * 'state -> Core.exp' * 'state}
+                  -> 'state -> Core.exp -> Core.exp * 'state
 end
 
 structure Decl : sig
diff --git a/src/core_util.sml b/src/core_util.sml
index 49182c09..df8465ae 100644
--- a/src/core_util.sml
+++ b/src/core_util.sml
@@ -578,6 +578,13 @@ fun exists {kind, con, exp} k =
         S.Return _ => true
       | S.Continue _ => false
 
+fun foldMap {kind, con, exp} s e =
+    case mapfold {kind = fn k => fn s => S.Continue (kind (k, s)),
+                  con = fn c => fn s => S.Continue (con (c, s)),
+                  exp = fn e => fn s => S.Continue (exp (e, s))} e s of
+        S.Continue v => v
+      | S.Return _ => raise Fail "CoreUtil.Exp.foldMap: Impossible"
+
 end
 
 structure Decl = struct
diff --git a/src/especialize.sig b/src/especialize.sig
new file mode 100644
index 00000000..df83e81b
--- /dev/null
+++ b/src/especialize.sig
@@ -0,0 +1,32 @@
+(* Copyright (c) 2008, Adam Chlipala
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * - Redistributions of source code must retain the above copyright notice,
+ *   this list of conditions and the following disclaimer.
+ * - Redistributions in binary form must reproduce the above copyright notice,
+ *   this list of conditions and the following disclaimer in the documentation
+ *   and/or other materials provided with the distribution.
+ * - The names of contributors may not be used to endorse or promote products
+ *   derived from this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
+ * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 
+ * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ * POSSIBILITY OF SUCH DAMAGE.
+ *)
+
+signature ESPECIALIZE = sig
+
+    val specialize : Core.file -> Core.file
+
+end
diff --git a/src/especialize.sml b/src/especialize.sml
new file mode 100644
index 00000000..a316ffaa
--- /dev/null
+++ b/src/especialize.sml
@@ -0,0 +1,176 @@
+(* Copyright (c) 2008, Adam Chlipala
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * - Redistributions of source code must retain the above copyright notice,
+ *   this list of conditions and the following disclaimer.
+ * - Redistributions in binary form must reproduce the above copyright notice,
+ *   this list of conditions and the following disclaimer in the documentation
+ *   and/or other materials provided with the distribution.
+ * - The names of contributors may not be used to endorse or promote products
+ *   derived from this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
+ * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 
+ * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ * POSSIBILITY OF SUCH DAMAGE.
+ *)
+
+structure ESpecialize :> ESPECIALIZE = struct
+
+open Core
+
+structure E = CoreEnv
+structure U = CoreUtil
+
+structure ILK = struct
+type ord_key = int list
+val compare = Order.joinL Int.compare
+end
+
+structure ILM = BinaryMapFn(ILK)
+structure IM = IntBinaryMap
+
+type func = {
+     name : string,
+     args : int ILM.map,
+     body : exp,
+     typ : con,
+     tag : string
+}
+
+type state = {
+     maxName : int,
+     funcs : func IM.map,
+     decls : (string * int * con * exp * string) list
+}
+
+fun kind (k, st) = (k, st)
+fun con (c, st) = (c, st)
+
+fun exp (e, st : state) =
+    let
+        fun getApp e =
+            case e of
+                ENamed f => SOME (f, [], [])
+              | EApp (e1, (ENamed x, _)) =>
+                (case getApp (#1 e1) of
+                     NONE => NONE
+                   | SOME (f, xs, xs') => SOME (f, xs @ [x], xs'))
+              | EApp (e1, e2) =>
+                (case getApp (#1 e1) of
+                     NONE => NONE
+                   | SOME (f, xs, xs') => SOME (f, xs, xs' @ [e2]))
+              | _ => NONE
+    in
+        case getApp e of
+            NONE => (e, st)
+          | SOME (_, [], _) => (e, st)
+          | SOME (f, xs, xs') =>
+            case IM.find (#funcs st, f) of
+                NONE => (e, st)
+              | SOME {name, args, body, typ, tag} =>
+                case ILM.find (args, xs) of
+                    SOME f' => (#1 (foldl (fn (e, arg) => (EApp (e, arg), ErrorMsg.dummySpan))
+                                          (ENamed f', ErrorMsg.dummySpan) xs'),
+                                st)
+                  | NONE =>
+                    let
+                        fun subBody (body, typ, xs) =
+                            case (#1 body, #1 typ, xs) of
+                                (_, _, []) => SOME (body, typ)
+                              | (EAbs (_, _, _, body'), TFun (_, typ'), x :: xs) =>
+                                subBody (E.subExpInExp (0, (ENamed x, ErrorMsg.dummySpan)) body',
+                                         typ',
+                                         xs)
+                              | _ => NONE
+                    in
+                        case subBody (body, typ, xs) of
+                            NONE => (e, st)
+                          | SOME (body', typ') =>
+                            let
+                                val f' = #maxName st
+                                val funcs = IM.insert (#funcs st, f, {name = name,
+                                                                      args = ILM.insert (args, xs, f'),
+                                                                      body = body,
+                                                                      typ = typ,
+                                                                      tag = tag})
+                                val st = {
+                                    maxName = f' + 1,
+                                    funcs = funcs,
+                                    decls = #decls st
+                                }
+
+                                val (body', st) = specExp st body'
+                                val e' = foldl (fn (e, arg) => (EApp (e, arg), ErrorMsg.dummySpan))
+                                               (ENamed f', ErrorMsg.dummySpan) xs'
+                            in
+                                (#1 e',
+                                 {maxName = #maxName st,
+                                  funcs = #funcs st,
+                                  decls = (name, f', typ', body', tag ^ "_espec") :: #decls st})
+                            end
+                    end
+    end
+
+and specExp st = U.Exp.foldMap {kind = kind, con = con, exp = exp} st
+
+fun decl (d, st) = (d, st)
+
+val specDecl = U.Decl.foldMap {kind = kind, con = con, exp = exp, decl = decl}
+
+fun specialize file =
+    let
+        fun doDecl (d, st) =
+            let
+                val (d', st) = specDecl st d
+
+                val funcs = #funcs st
+                val funcs =
+                    case #1 d of
+                        DVal (x, n, c, e as (EAbs _, _), tag) =>
+                        IM.insert (funcs, n, {name = x,
+                                              args = ILM.empty,
+                                              body = e,
+                                              typ = c,
+                                              tag = tag})
+                      | DValRec vis =>
+                        foldl (fn ((x, n, c, e, tag), funcs) =>
+                                  IM.insert (funcs, n, {name = x,
+                                                        args = ILM.empty,
+                                                        body = e,
+                                                        typ = c,
+                                                        tag = tag}))
+                              funcs vis
+                      | _ => funcs
+
+                val ds =
+                    case #decls st of
+                        [] => [d']
+                      | vis => [(DValRec vis, ErrorMsg.dummySpan), d']
+            in
+                (ds, {maxName = #maxName st,
+                      funcs = funcs,
+                      decls = []})
+            end
+
+        val (ds, _) = ListUtil.foldlMapConcat doDecl
+                      {maxName = U.File.maxName file + 1,
+                       funcs = IM.empty,
+                       decls = []}
+                      file
+    in
+        ds
+    end
+
+
+end
diff --git a/src/reduce.sml b/src/reduce.sml
index 927c8ff1..8dc4527f 100644
--- a/src/reduce.sml
+++ b/src/reduce.sml
@@ -37,36 +37,8 @@ structure U = CoreUtil
 val liftConInCon = E.liftConInCon
 val subConInCon = E.subConInCon
 val liftConInExp = E.liftConInExp
-
-val liftExpInExp =
-    U.Exp.mapB {kind = fn k => k,
-                con = fn _ => fn c => c,
-                exp = fn bound => fn e =>
-                                     case e of
-                                         ERel xn =>
-                                         if xn < bound then
-                                             e
-                                         else
-                                             ERel (xn + 1)
-                                       | _ => e,
-                bind = fn (bound, U.Exp.RelE _) => bound + 1
-                        | (bound, _) => bound}
-
-val subExpInExp =
-    U.Exp.mapB {kind = fn k => k,
-                con = fn _ => fn c => c,
-                exp = fn (xn, rep) => fn e =>
-                                  case e of
-                                      ERel xn' =>
-                                      (case Int.compare (xn', xn) of
-                                           EQUAL => #1 rep
-                                         | GREATER=> ERel (xn' - 1)
-                                         | LESS => e)
-                                    | _ => e,
-                bind = fn ((xn, rep), U.Exp.RelE _) => (xn+1, liftExpInExp 0 rep)
-                        | ((xn, rep), U.Exp.RelC _) => (xn, liftConInExp 0 rep)
-                        | (ctx, _) => ctx}
-
+val liftExpInExp = E.liftExpInExp
+val subExpInExp = E.subExpInExp
 val liftConInExp = E.liftConInExp
 val subConInExp = E.subConInExp
 
diff --git a/src/sources b/src/sources
index 3568279c..ebf71d9e 100644
--- a/src/sources
+++ b/src/sources
@@ -93,6 +93,9 @@ unpoly.sml
 specialize.sig
 specialize.sml
 
+especialize.sig
+especialize.sml
+
 tag.sig
 tag.sml