Fix case on datatype.

5 files changed, 116 insertions, 55 deletions

diff --git a/cps.sml b/cps.sml
index 894c57e..83d07f7 100644
--- a/cps.sml
+++ b/cps.sml
@@ -67,23 +67,26 @@ struct
           val contAddr = Gensym.new ()
           val otherwiseAddr = Gensym.new ()
           val arg = Gensym.new ()
-          fun go _ [] cont = toCPS otherwise cont
+          fun go _ [] cont = Syntax.CApp (Syntax.VVar otherwiseAddr, [])
             | go v [(x, arm)] cont =
-                let val b = Gensym.new ()
-                in
-                  Syntax.CPrimop
-                    ( Syntax.PEq
-                    , [v, Syntax.VInt x]
-                    , [b]
-                    , [ Syntax.CPrimop
-                          ( Syntax.PIf
-                          , [Syntax.VVar b]
-                          , []
-                          , [toCPS arm cont, Syntax.CApp (Syntax.VVar otherwiseAddr, [])]
-                          )
-                      ]
-                    )
-                end
+                (case otherwise of
+                  NONE => toCPS arm cont
+                | SOME _ =>
+                    let val b = Gensym.new ()
+                    in
+                      Syntax.CPrimop
+                        ( Syntax.PEq
+                        , [v, Syntax.VInt x]
+                        , [b]
+                        , [ Syntax.CPrimop
+                              ( Syntax.PIf
+                              , [Syntax.VVar b]
+                              , []
+                              , [toCPS arm cont, Syntax.CApp (Syntax.VVar otherwiseAddr, [])]
+                              )
+                          ]
+                        )
+                    end)
             | go v arms cont =
                 let
                   val b = Gensym.new ()
@@ -106,11 +109,14 @@ struct
                     )
                 end
           fun contFunc x = Syntax.CApp (Syntax.VVar contAddr, [x])
+          val fixFuncs = [(contAddr, [arg], cont (Syntax.VVar arg))]
+          val fixFuncs =
+            case otherwise of
+              NONE => fixFuncs
+            | SOME otherwise => (otherwiseAddr, [], toCPS otherwise contFunc) :: fixFuncs
         in
           Syntax.CFix
-            ( [ (contAddr, [arg], cont (Syntax.VVar arg))
-              , (otherwiseAddr, [], toCPS otherwise contFunc)
-              ]
+            ( fixFuncs
             , toCPS expr (fn v => go v sortedArms contFunc)
             )
         end
diff --git a/elab.sml b/elab.sml
index 3a361a6..eb168c5 100644
--- a/elab.sml
+++ b/elab.sml
@@ -12,12 +12,12 @@ struct
     | "div" => Syntax.PDiv
     | _ => raise Fail ("invalid op: " ^ s)
 
-  fun mapIdx (f : int * 'a -> 'b) (l : 'a list) : 'b list =
-    let
-      fun go _ [] = []
-        | go i (x :: xs) = f (i, x) :: go (i + 1) xs
-    in go 0 l
-    end
+  fun enumerate (l : 'a list) : (int * 'a) list =
+  let
+    fun go _ [] = []
+      | go i (x :: xs) = (i, x) :: go (i + 1) xs
+  in go 0 l
+  end
 
   fun hdstls ([] : 'a list list) : ('a list * 'a list list) option = SOME ([], [])
     | hdstls ([] :: _) = NONE
@@ -32,7 +32,7 @@ struct
         NONE => []
       | SOME (heads, tails) => heads :: transpose tails
 
-  type env = int StringMap.map * int StringMap.map
+  type env = int StringMap.map * (int * int) StringMap.map
 
   fun bindVar (name : string) (sym : int) ((vars, types) : env) : env = (StringMap.insert name sym vars, types)
 
@@ -43,16 +43,23 @@ struct
 
   fun bindDataCons (cons : (string * Syntax.etype option) list) ((vars, types) : env) : env =
   let
+    val nCons = length cons
     val (_, vars, types) =
       foldl
         (fn ((name, _), (i, vars, types)) =>
-          (i + 1, StringMap.insert name (Gensym.new ()) vars, StringMap.insert name i types))
+          (i + 1, StringMap.insert name (Gensym.new ()) vars, StringMap.insert name (i, nCons) types))
         (0, vars, types)
         cons
   in (vars, types)
   end
 
-  fun lookupCon (name : string) ((_, types) : env) : int option = StringMap.lookup name types
+  fun lookupCon (name : string) ((_, types) : env) : int option =
+    Option.map (fn (i, _) => i) (StringMap.lookup name types)
+
+  fun nConstructors (name : string) ((_, types) : env) : int =
+    let val (_, n) = valOf (StringMap.lookup name types)
+    in n
+    end
 
   fun patternMatrix (arms : Syntax.pat list) : Syntax.pat list list =
   let val ts =
@@ -92,12 +99,12 @@ struct
   in
     case cols of
       [] => [expr]
-    | _ => List.concat (mapIdx (fn (i, col) => occurrenceVector (Syntax.LSelect (i, expr)) col) cols)
+    | _ => List.concat (map (fn (i, col) => occurrenceVector (Syntax.LSelect (i, expr)) col) (enumerate cols))
   end
 
   fun patternBindings (expr : Syntax.lexp) (Syntax.PVar v) : (string * Syntax.lexp) list = [(v, expr)]
     | patternBindings expr (Syntax.PTuple t) =
-        List.concat (mapIdx (fn (i, p) => patternBindings (Syntax.LSelect (i, expr)) p) t)
+        List.concat (map (fn (i, p) => patternBindings (Syntax.LSelect (i, expr)) p) (enumerate t))
     | patternBindings expr (Syntax.PCon (_, arg)) = patternBindings (Syntax.LSelect (1, expr)) arg
     | patternBindings _ _ = []
 
@@ -113,6 +120,22 @@ struct
 
   fun specialize (env : env) (n : int) (patterns : Syntax.pat list list, occurrences : Syntax.lexp list, actions : Syntax.lexp list) : Syntax.pat list list * Syntax.lexp list * Syntax.lexp list =
     let
+      val con1 =
+        List.find
+          (fn Syntax.PCon (name, _) => valOf (lookupCon name env) = n
+            | _ => false)
+          (map hd patterns)
+      val newTupleSize =
+        case con1 of
+          SOME (Syntax.PCon (_, Syntax.PTuple t)) => length t
+        | _ => 0
+      val occHead = hd occurrences
+      val occRest = tl occurrences
+      val occurrences =
+        if newTupleSize = 0
+        then Syntax.LSelect (1, occHead) :: occRest
+        else
+          List.tabulate (newTupleSize, fn i => Syntax.LSelect (i, Syntax.LSelect (1, occHead))) @ occRest
       fun specializeRow (Syntax.PInt i :: rest) =
             if i = n then SOME (Syntax.PWild :: rest) else NONE
         | specializeRow (Syntax.PWild :: rest) = SOME (Syntax.PWild :: rest)
@@ -131,20 +154,20 @@ struct
           (List.mapPartial
             (fn (p, a) =>
               Option.map (fn p => (p, a)) (specializeRow p))
-            (ListPair.zip (patterns, actions)))
-    in (patterns, tl occurrences, actions)
+            (ListPair.zipEq (patterns, actions)))
+    in (patterns, occurrences, actions)
     end
 
   fun default (patterns : Syntax.pat list list, occurrences : Syntax.lexp list, actions : Syntax.lexp list) : Syntax.pat list list * Syntax.lexp list * Syntax.lexp list =
     let
       val (patterns, actions) =
         ListPair.unzip
-          (List.mapPartial
-            (fn (Syntax.PWild :: rest, a) => SOME (rest, a)
-              | (Syntax.PVar _ :: rest, a) => SOME (rest, a)
-              | _ => NONE)
-            (ListPair.zip (patterns, actions)))
-    in (patterns, tl occurrences, actions)
+          (List.filter
+            (fn (Syntax.PWild :: _, _) => true
+              | (Syntax.PVar _ :: _, _) => true
+              | _ => false)
+            (ListPair.zipEq (patterns, actions)))
+    in (patterns, occurrences, actions)
     end
 
   (* There's ambiguity between pattern variables and constructors that can only
@@ -165,8 +188,9 @@ struct
         let val refutablePattern =
           List.find
             (fn (_, Syntax.PInt _) => true
+              | (_, Syntax.PCon _) => true
               | _ => false)
-            (mapIdx (fn x => x) firstRow)
+            (enumerate firstRow)
         in
           case refutablePattern of
             NONE => hd actions
@@ -176,21 +200,35 @@ struct
                   if i = 0
                   then (patterns, occurrences)
                   else swap i patterns occurrences
+                val firstCol = map hd patterns
                 val signatures =
                   map (fn (x, _) => x)
                     (IntMap.toList
                       (foldl
                         (fn (Syntax.PInt i, acc) => IntMap.insert i true acc
+                          | (Syntax.PCon (c, _), acc) => IntMap.insert (valOf (lookupCon c env)) true acc
                           | (_, acc) => acc)
                         IntMap.empty
-                        (map hd patterns)))
+                        firstCol))
+                val nCons =
+                  case List.find (fn (Syntax.PCon _) => true | _ => false) firstCol of
+                    SOME (Syntax.PCon (name, _)) => nConstructors name env
+                  | _ => ~1
+                val defaultCase =
+                  if length signatures = nCons
+                  then NONE
+                  else SOME (compilePatternMatching env (default (patterns, occurrences, actions)))
+                val switchOperand =
+                  if nCons < 0
+                  then hd occurrences
+                  else Syntax.LSelect (0, hd occurrences)
               in
                 Syntax.LSwitch
-                  ( hd occurrences
+                  ( switchOperand
                   , map
                       (fn i => (i, compilePatternMatching env (specialize env i (patterns, occurrences, actions))))
                       signatures
-                  , compilePatternMatching env (default (patterns, occurrences, actions))
+                  , defaultCase
                   )
               end
         end
@@ -258,13 +296,25 @@ struct
     | Syntax.ELet (Syntax.DDatatype (name, cons) :: decls, body) =>
         let
           val env = bindDataCons cons env
-          fun go _ [] = []
-            | go i ((name, _) :: cons) =
-                let val v = Gensym.new ()
-                in (lookupVar name env, v, Syntax.LRecord [Syntax.LInt i, Syntax.LVar v]) :: go (i + 1) cons
-                end
+          val funs =
+            List.mapPartial
+              (fn (_, (_, NONE)) => NONE
+                | (i, (name, _)) =>
+                    let val v = Gensym.new ()
+                    in SOME (lookupVar name env, v, Syntax.LRecord [Syntax.LInt i, Syntax.LVar v])
+                    end)
+              (enumerate cons)
+          val vals =
+            List.mapPartial
+              (fn (i, (name, NONE)) => SOME (lookupVar name env, Syntax.LRecord [Syntax.LInt i])
+                | _ => NONE)
+              (enumerate cons)
         in
-          Syntax.LFix (go 0 cons, elab env (Syntax.ELet (decls, body)))
+          foldl
+            (fn ((v, x), acc) =>
+              Syntax.LApp (Syntax.LFn (v, acc), x))
+            (Syntax.LFix (funs, elab env (Syntax.ELet (decls, body))))
+            vals
         end
     | Syntax.ELet (Syntax.DVal (pat, v) :: decls, body) =>
         elab env (Syntax.ECase (v, [(pat, Syntax.ELet (decls, body))]))
diff --git a/syntax.sml b/syntax.sml
index 767851d..aa94b09 100644
--- a/syntax.sml
+++ b/syntax.sml
@@ -57,7 +57,7 @@ struct
   | LRecord of lexp list
   | LSelect of int * lexp
   | LPrim of primop
-  | LSwitch of lexp * (int * lexp) list * lexp
+  | LSwitch of lexp * (int * lexp) list * lexp option
 
   (* CPS *)
   datatype value =
@@ -173,7 +173,7 @@ struct
     | LRecord l => "LRecord " ^ listToString lexpToString l
     | LSelect (i, r) => "LSelect (" ^ Int.toString i ^ ", " ^ lexpToString r ^ ")"
     | LPrim p => "LPrim " ^ primopToString p
-    | LSwitch (e, arms, otherwise) => "LSwitch (" ^ lexpToString e ^ ", " ^ listToString (fn (x, e) => "(" ^ Int.toString x ^ ", " ^ lexpToString e ^ ")") arms ^ ", " ^ lexpToString otherwise ^ ")"
+    | LSwitch (e, arms, otherwise) => "LSwitch (" ^ lexpToString e ^ ", " ^ listToString (fn (x, e) => "(" ^ Int.toString x ^ ", " ^ lexpToString e ^ ")") arms ^ ", " ^ optionToString lexpToString otherwise ^ ")"
 
   fun valueToString (x : value) : string =
     case x of
diff --git a/tests/16-datatype.sml b/tests/16-datatype.sml
index ee4dd16..7514606 100644
--- a/tests/16-datatype.sml
+++ b/tests/16-datatype.sml
@@ -1,6 +1,5 @@
-datatype D = A | B of int
+datatype D = D of int
 
-val _ =
-  case B 42 of
-    B x => __builtin "exit" x
-  | _ => __builtin "exit" 0
+fun f (D x) = __builtin "exit" x
+
+val _ = f (D 42)
diff --git a/tests/17-case-datatype.sml b/tests/17-case-datatype.sml
new file mode 100644
index 0000000..974f263
--- /dev/null
+++ b/tests/17-case-datatype.sml
@@ -0,0 +1,6 @@
+datatype D = A | B
+
+val _ =
+  case B of
+    A => __builtin "exit" 0
+  | B => __builtin "exit" 42