amelia
/
ahc-old


module M = import "data/map.ml"

module S = import "data/set.ml"


open import "prelude.ml"

open import "./lang.ml"

open import "./lib/monads.ml"


type addr =

 | Combinator of string

 | Arg of int

 | Int of int


type gm_code =

 | Push of addr

 | Update of int

 | Pop of int

 | Unwind

 | Mkap

 | Add | Sub | Mul | Div | Eval

 | Iszero of list gm_code * list gm_code


instance show gm_code begin

 let show = function

 | Mkap -> "Mkap"

 | Unwind -> "Unwind"

 | Push (Combinator k) -> "Push " ^ k

 | Push (Arg i) -> "Arg " ^ show i

 | Push (Int i) -> "Int " ^ show i

 | Update n -> "Update " ^ show n

 | Pop n -> "Pop " ^ show n

 | Add -> "Add"

 | Mul -> "Mul"

 | Sub -> "Sub"

 | Div -> "Div"

 | Eval -> "Eval"

 | Iszero p -> "Iszero " ^ show p

end


type program_item =

 | Sc of string * int * list gm_code

 | Fd of fdecl


let rec lambda_lift = function

 | Ref v -> pure (Ref v)

 | Lit v -> pure (Lit v)

 | App (f, x) -> (| app (lambda_lift f) (lambda_lift x) |)

 | Lam (v, x) ->

 let! body = lambda_lift x

 let! (i, defs, known_sc) = get


 let vars =

 x |> free_vars

 |> S.delete v

 |> flip S.difference known_sc

 |> S.members


 let def = ("Lam" ^ show i, vars ++ [v], body)

 let app = foldl (fun f -> app f # Ref) (Ref ("Lam" ^ show i)) vars


 put (i + 1, Decl def :: defs, known_sc)

 |> map (const app)

 | Case (sc, alts) ->

 alts

 |> map (fun (_, x) -> x)

 |> foldl app sc

 |> lambda_lift


let rec eta_contract = function

 | Decl (n, a, e) as dec ->

 match a, e with

 | [], _ -> dec

 | xs, App (f, Ref v) ->

 if v == last xs && not (S.member v (free_vars f)) then

 eta_contract (Decl (n, init a, f))

 else

 dec

 | _, _ -> dec

 | Data c -> Data c

 | Foreign i -> Foreign i


let rec lambda_lift_sc = function

 | Decl (n, a, e) ->

 match e with

 | Lam (v, e) -> lambda_lift_sc (Decl (n, a ++ [v], e))

 | _ ->

 let! e = lambda_lift e

 let! _ = modify (fun (a, b, s) -> (a, b, S.insert n s))

 pure (Decl (n, a, e))

 | Data c -> Data c |> pure

 | Foreign i -> Foreign i |> pure


type dlist 'a <- list 'a -> list 'a


let cg_prim (Fimport { var, fent }) =

 let prim_math_op x =

 [ Push (Arg 0), Eval, Push (Arg 2), Eval, x, Update 2, Pop 2, Unwind ]

 let prim_equality =

 [ Push (Arg 0), Eval (* x, arg0, arg1, arg2, arg3 *)

 , Push (Arg 2), Eval (* y, x, arg0, arg1, arg2, arg3 *)

 , Sub (* y - x, arg0, arg1, arg2, arg3 *)

 , Iszero ([ Push (Arg 3) ], [ Push (Arg 4) ])

 , Push (Int 0), Mkap, Update 4, Pop 4, Unwind ]

 match fent with

 | "add" -> Sc (var, 2, prim_math_op Add)

 | "sub" -> Sc (var, 2, prim_math_op Sub)

 | "mul" -> Sc (var, 2, prim_math_op Mul)

 | "div" -> Sc (var, 2, prim_math_op Div)

 | "equ" -> Sc (var, 4, prim_equality)

 | e -> error @@ "No such primitive " ^ e


let rec compile (env : M.t string int) = function

 | Ref v ->

 match M.lookup v env with

 | Some i -> (Push (Arg i) ::)

 | None -> (Push (Combinator v) ::)


 | App (f, x) ->

 let f = compile env f

 let x = compile (map (1 +) env) x

 f # x # (Mkap ::)


 | Lam _ ->

 error "Can not compile lambda expression, did you forget to lift?"

 | Case _ ->

 error "Can not compile case expression, did you forget to lift?"

 | Lit i -> (Push (Int i) ::)


let supercomb (_, args, exp) =

 let env = M.from_list (zip args [0..length args])

 let k = compile (M.from_list (zip args [0..length args])) exp

 k [Update (length env), Pop (length env), Unwind]


let known_scs = S.from_list [ "getchar", "putchar" ]


let program decs =

 let (decs, (_, lams, _)) =

 run_state (traverse (lambda_lift_sc # eta_contract) decs) (0, [], known_scs)

 let define nm =

 let! x = get

 if nm `S.member` x then

 error @@ "Redefinition of value " ^ nm

 else

 modify (S.insert nm)


 let go =

 flip traverse (lams ++ decs) @@ function

 | Decl ((nm, args, _) as sc) ->

 let! _ = define nm

 let code = supercomb sc

 Sc (nm, length args, code) |> pure

 | Data _ -> error "data declaration in compiler"

 | Foreign (Fimport { cc = Prim, var } as fi) ->

 let! _ = define var

 pure (cg_prim fi)

 | Foreign f -> pure (Fd f)

 let (out, _) = run_state go S.empty

 out