From bac5ff10422778a32a32a1b285eb170f60036e75 Mon Sep 17 00:00:00 2001
From: Dan Doel <dan.doel@gmail.com>
Date: Thu, 18 Jun 2020 16:44:53 -0400
Subject: [PATCH] Eta expand constructors before floating

---
 parser-typechecker/src/Unison/Runtime/ANF.hs  | 24 +++++++++++++++++++
 .../src/Unison/Runtime/Interface.hs           |  6 +++++
 unison-core/src/Unison/Var.hs                 |  3 +++
 3 files changed, 33 insertions(+)

diff --git a/parser-typechecker/src/Unison/Runtime/ANF.hs b/parser-typechecker/src/Unison/Runtime/ANF.hs
index fd9808161..bc92e4131 100644
--- a/parser-typechecker/src/Unison/Runtime/ANF.hs
+++ b/parser-typechecker/src/Unison/Runtime/ANF.hs
@@ -41,6 +41,7 @@ module Unison.Runtime.ANF
   , POp(..)
   , IOp(..)
   , close
+  , constructorEta
   , float
   , lamLift
   , ANormalBF(..)
@@ -232,6 +233,29 @@ deannotate = ABT.visitPure $ \case
 lamLift :: (Var v, Monoid a) => Term v a -> Term v a
 lamLift = float . close Set.empty . deannotate
 
+constructorEta
+  :: (Var v, Monoid a)
+  => Map (Reference,Int) Int -> Term v a -> Term v a
+constructorEta dat = ABT.visitPure $ \case
+  Apps' f@(Constructor' r t) args -> eta r t f args
+  Apps' f@(Request' r t) args -> eta r t f args
+  f@(Constructor' r t) -> eta r t f []
+  f@(Request' r t) -> eta r t f []
+  _ -> Nothing
+  where
+  frsh avoid _ =
+    let v = Var.freshIn avoid $ typed Var.Eta
+    in (Set.insert v avoid, v)
+  eta r t f args
+    | Just n <- Map.lookup (r,t) dat
+    , vs <- snd . mapAccumL frsh fvs $ drop (length args) [1..n]
+    , nargs <- var mempty <$> vs
+    = Just . lam' mempty vs . apps' f $ args' ++ nargs
+    | otherwise = Just (apps' f args')
+    where
+    fvs = foldMap freeVars args
+    args' = constructorEta dat <$> args
+
 optimize :: forall a v . (Semigroup a, Var v) => Term v a -> Term v a
 optimize t = go t where
   ann = ABT.annotation
diff --git a/parser-typechecker/src/Unison/Runtime/Interface.hs b/parser-typechecker/src/Unison/Runtime/Interface.hs
index 51f6962df..51252c25b 100644
--- a/parser-typechecker/src/Unison/Runtime/Interface.hs
+++ b/parser-typechecker/src/Unison/Runtime/Interface.hs
@@ -51,6 +51,11 @@ data EvalCtx v
   , backrefTm :: EnumMap Word64 (Term v)
   }
 
+uncurryDspec :: DataSpec -> Map.Map (Reference,Int) Int
+uncurryDspec = Map.fromList . concatMap f . Map.toList
+  where
+  f (r,l) = zipWith (\n c -> ((r,n),c)) [0..] $ either id id l
+
 numberLetRec :: Word64 -> Term v -> EnumMap Word64 (Term v)
 numberLetRec frsh (Tm.LetRecNamed' bs e)
   = mapFromList . zip [frsh..] $ e : map snd bs
@@ -150,6 +155,7 @@ compileTerm w tm ctx
   . superNormalize (ref $ refTm ctx) (ref $ refTy ctx)
   . lamLift
   . splitPatterns (dspec ctx)
+  . constructorEta (uncurryDspec $ dspec ctx)
   $ tm
   where
   frsh = freshTm ctx
diff --git a/unison-core/src/Unison/Var.hs b/unison-core/src/Unison/Var.hs
index 1f9584a6d..260f2204a 100644
--- a/unison-core/src/Unison/Var.hs
+++ b/unison-core/src/Unison/Var.hs
@@ -52,6 +52,7 @@ rawName typ = case typ of
   Inference TypeConstructorArg -> "𝕦"
   MissingResult -> "_"
   Blank -> "_"
+  Eta -> "_eta"
   ANFBlank -> "_anf"
   Float -> "_float"
   Pattern -> "_pattern"
@@ -104,6 +105,8 @@ data Type
   --  > 1 + 1
   --    has kind ""
   | UnnamedWatch WatchKind Text -- guid
+  -- An unnamed variable for constructor eta expansion
+  | Eta
   -- An unnamed variable introduced by ANF transformation
   | ANFBlank
   -- An unnamed variable for a floated lambda