adamc@443: (* Copyright (c) 2008, Adam Chlipala
adamc@443:  * All rights reserved.
adamc@443:  *
adamc@443:  * Redistribution and use in source and binary forms, with or without
adamc@443:  * modification, are permitted provided that the following conditions are met:
adamc@443:  *
adamc@443:  * - Redistributions of source code must retain the above copyright notice,
adamc@443:  *   this list of conditions and the following disclaimer.
adamc@443:  * - Redistributions in binary form must reproduce the above copyright notice,
adamc@443:  *   this list of conditions and the following disclaimer in the documentation
adamc@443:  *   and/or other materials provided with the distribution.
adamc@443:  * - The names of contributors may not be used to endorse or promote products
adamc@443:  *   derived from this software without specific prior written permission.
adamc@443:  *
adamc@443:  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
adamc@443:  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
adamc@443:  * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
adamc@443:  * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
adamc@443:  * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 
adamc@443:  * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
adamc@443:  * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
adamc@443:  * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
adamc@443:  * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
adamc@443:  * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
adamc@443:  * POSSIBILITY OF SUCH DAMAGE.
adamc@443:  *)
adamc@443: 
adamc@443: structure ESpecialize :> ESPECIALIZE = struct
adamc@443: 
adamc@443: open Core
adamc@443: 
adamc@443: structure E = CoreEnv
adamc@443: structure U = CoreUtil
adamc@443: 
adamc@443: structure ILK = struct
adamc@443: type ord_key = int list
adamc@443: val compare = Order.joinL Int.compare
adamc@443: end
adamc@443: 
adamc@443: structure ILM = BinaryMapFn(ILK)
adamc@443: structure IM = IntBinaryMap
adamc@443: 
adamc@443: type func = {
adamc@443:      name : string,
adamc@443:      args : int ILM.map,
adamc@443:      body : exp,
adamc@443:      typ : con,
adamc@443:      tag : string
adamc@443: }
adamc@443: 
adamc@443: type state = {
adamc@443:      maxName : int,
adamc@443:      funcs : func IM.map,
adamc@443:      decls : (string * int * con * exp * string) list
adamc@443: }
adamc@443: 
adamc@443: fun kind (k, st) = (k, st)
adamc@443: fun con (c, st) = (c, st)
adamc@443: 
adamc@443: fun exp (e, st : state) =
adamc@443:     let
adamc@443:         fun getApp e =
adamc@443:             case e of
adamc@443:                 ENamed f => SOME (f, [], [])
adamc@443:               | EApp (e1, (ENamed x, _)) =>
adamc@443:                 (case getApp (#1 e1) of
adamc@443:                      NONE => NONE
adamc@443:                    | SOME (f, xs, xs') => SOME (f, xs @ [x], xs'))
adamc@443:               | EApp (e1, e2) =>
adamc@443:                 (case getApp (#1 e1) of
adamc@443:                      NONE => NONE
adamc@443:                    | SOME (f, xs, xs') => SOME (f, xs, xs' @ [e2]))
adamc@443:               | _ => NONE
adamc@443:     in
adamc@443:         case getApp e of
adamc@443:             NONE => (e, st)
adamc@443:           | SOME (_, [], _) => (e, st)
adamc@443:           | SOME (f, xs, xs') =>
adamc@443:             case IM.find (#funcs st, f) of
adamc@443:                 NONE => (e, st)
adamc@443:               | SOME {name, args, body, typ, tag} =>
adamc@443:                 case ILM.find (args, xs) of
adamc@443:                     SOME f' => (#1 (foldl (fn (e, arg) => (EApp (e, arg), ErrorMsg.dummySpan))
adamc@443:                                           (ENamed f', ErrorMsg.dummySpan) xs'),
adamc@443:                                 st)
adamc@443:                   | NONE =>
adamc@443:                     let
adamc@443:                         fun subBody (body, typ, xs) =
adamc@443:                             case (#1 body, #1 typ, xs) of
adamc@443:                                 (_, _, []) => SOME (body, typ)
adamc@443:                               | (EAbs (_, _, _, body'), TFun (_, typ'), x :: xs) =>
adamc@443:                                 subBody (E.subExpInExp (0, (ENamed x, ErrorMsg.dummySpan)) body',
adamc@443:                                          typ',
adamc@443:                                          xs)
adamc@443:                               | _ => NONE
adamc@443:                     in
adamc@443:                         case subBody (body, typ, xs) of
adamc@443:                             NONE => (e, st)
adamc@443:                           | SOME (body', typ') =>
adamc@443:                             let
adamc@443:                                 val f' = #maxName st
adamc@443:                                 val funcs = IM.insert (#funcs st, f, {name = name,
adamc@443:                                                                       args = ILM.insert (args, xs, f'),
adamc@443:                                                                       body = body,
adamc@443:                                                                       typ = typ,
adamc@443:                                                                       tag = tag})
adamc@443:                                 val st = {
adamc@443:                                     maxName = f' + 1,
adamc@443:                                     funcs = funcs,
adamc@443:                                     decls = #decls st
adamc@443:                                 }
adamc@443: 
adamc@443:                                 val (body', st) = specExp st body'
adamc@443:                                 val e' = foldl (fn (e, arg) => (EApp (e, arg), ErrorMsg.dummySpan))
adamc@443:                                                (ENamed f', ErrorMsg.dummySpan) xs'
adamc@443:                             in
adamc@443:                                 (#1 e',
adamc@443:                                  {maxName = #maxName st,
adamc@443:                                   funcs = #funcs st,
adamc@443:                                   decls = (name, f', typ', body', tag ^ "_espec") :: #decls st})
adamc@443:                             end
adamc@443:                     end
adamc@443:     end
adamc@443: 
adamc@443: and specExp st = U.Exp.foldMap {kind = kind, con = con, exp = exp} st
adamc@443: 
adamc@443: fun decl (d, st) = (d, st)
adamc@443: 
adamc@443: val specDecl = U.Decl.foldMap {kind = kind, con = con, exp = exp, decl = decl}
adamc@443: 
adamc@443: fun specialize file =
adamc@443:     let
adamc@443:         fun doDecl (d, st) =
adamc@443:             let
adamc@443:                 val (d', st) = specDecl st d
adamc@443: 
adamc@443:                 val funcs = #funcs st
adamc@443:                 val funcs =
adamc@443:                     case #1 d of
adamc@443:                         DVal (x, n, c, e as (EAbs _, _), tag) =>
adamc@443:                         IM.insert (funcs, n, {name = x,
adamc@443:                                               args = ILM.empty,
adamc@443:                                               body = e,
adamc@443:                                               typ = c,
adamc@443:                                               tag = tag})
adamc@443:                       | DValRec vis =>
adamc@443:                         foldl (fn ((x, n, c, e, tag), funcs) =>
adamc@443:                                   IM.insert (funcs, n, {name = x,
adamc@443:                                                         args = ILM.empty,
adamc@443:                                                         body = e,
adamc@443:                                                         typ = c,
adamc@443:                                                         tag = tag}))
adamc@443:                               funcs vis
adamc@443:                       | _ => funcs
adamc@443: 
adamc@443:                 val ds =
adamc@443:                     case #decls st of
adamc@443:                         [] => [d']
adamc@443:                       | vis => [(DValRec vis, ErrorMsg.dummySpan), d']
adamc@443:             in
adamc@443:                 (ds, {maxName = #maxName st,
adamc@443:                       funcs = funcs,
adamc@443:                       decls = []})
adamc@443:             end
adamc@443: 
adamc@443:         val (ds, _) = ListUtil.foldlMapConcat doDecl
adamc@443:                       {maxName = U.File.maxName file + 1,
adamc@443:                        funcs = IM.empty,
adamc@443:                        decls = []}
adamc@443:                       file
adamc@443:     in
adamc@443:         ds
adamc@443:     end
adamc@443: 
adamc@443: 
adamc@443: end