more erasure

lib/Quox/Untyped/Erase.idr

@@ -1,14 +1,15 @@
 module Quox.Untyped.Erase
 
 import Quox.Definition as Q
+import Quox.Pretty
 import Quox.Syntax.Term.Base as Q
 import Quox.Syntax.Term.Subst
-import Quox.Untyped.Syntax as U
 import Quox.Typing
+import Quox.Untyped.Syntax as U
 import Quox.Whnf
-import Quox.Pretty
 
 import Quox.EffExtra
+import Data.List1
 import Data.Singleton
 import Data.SnocVect
 import Language.Reflection
@@ -28,12 +29,6 @@ isErased Zero = Erased
 isErased One  = Kept
 isErased Any  = Kept
 
-public export
-ifErased : Qty -> Lazy a -> Lazy a -> a
-ifErased pi x y = case isErased pi of
-  Erased => x
-  Kept   => y
-
 
 public export
 ErasureContext : Nat -> Nat -> Type
@@ -98,23 +93,18 @@ export covering
 computeElimType : ErasureContext d n -> SQty -> Elim d n -> Eff Erase (Term d n)
 computeElimType ctx sg e = do
   defs <- askAt DEFS
-  liftWhnf $ do
   let ctx = toWhnfContext ctx
-    Element e enf <- whnf defs ctx sg e
+  liftWhnf $ do
+    Element e _ <- whnf defs ctx sg e
     computeElimType defs ctx sg e
 
 
 private %macro
 wrapExpect : TTImp ->
              Elab (TyContext d n -> Loc -> Term d n -> Eff Erase a)
-wrapExpect f {a} = do
+wrapExpect f_ = do
-  f' <- check `(\x => ~(f) x)
+  f <- check `(\x => ~(f_) x)
-  pure $ \ctx, loc, s => wrapExpect' f' ctx loc s
+  pure $ \ctx, loc, s => liftWhnf $ f !(askAt DEFS) ctx SZero loc s
-where
-  wrapExpect' : (Q.Definitions -> TyContext d n -> SQty -> Loc -> Term d n ->
-                    Eff [Except TypeError, NameGen] a) ->
-                TyContext d n -> Loc -> Term d n -> Eff Erase a
-  wrapExpect' f ctx loc s = liftWhnf $ f !(askAt DEFS) ctx SZero loc s
 
 
 public export
@@ -263,22 +253,42 @@ eraseElim ctx (App fun arg loc) = do
     Kept   => do arg <- eraseTerm ctx targ arg
                  pure $ EraRes ty $ App efun.term arg loc
 
---                     e ⤋ e' ⇒ (x : A) × B
+--                        e ⇒ (x : A) × B
 --    x : A, y : B | ρ.x, ρ.y ⊢ s ⤋ s' ⇐ R[((x,y) ∷ (x : A) × B)/z]
---  x̃ ≔ if ρ = 0 then ⌷ else fst e'   ỹ ≔ if ρ = 0 then ⌷ else snd e'
 -- -------------------------------------------------------------------
---   (caseρ e return z ⇒ R of {(x, y) ⇒ s}) ⤋ s'[x̃/x, ỹ/y] ⇒ R[e/z]
+--   (case0 e return z ⇒ R of {(x, y) ⇒ s}) ⤋ s'[⌷/x, ⌷/y] ⇒ R[e/z]
+--
+--                     e ⤋ e' ⇒ (x : A) × B   ρ ≠ 0
+--    x : A, y : B | ρ.x, ρ.y ⊢ s ⤋ s' ⇐ R[((x,y) ∷ (x : A) × B)/z]
+-- ----------------------------------------------------------------------------
+--   (caseρ e return z ⇒ R of {(x, y) ⇒ s}) ⤋
+--     ⤋
+--   let xy = e' in let x = fst xy in let y = snd xy in s' ⇒ R[e/z]
 eraseElim ctx (CasePair qty pair ret body loc) = do
-  epair <- eraseElim ctx pair
-  let ty = sub1 (ret // shift 2) $
-           Ann (Pair (BVT 0 loc) (BVT 1 loc) loc) (weakT 2 epair.type) loc
-  (tfst, tsnd) <- wrapExpect `(expectSig) ctx loc epair.type
   let [< x, y] = body.names
+  case isErased qty of
+    Kept => do
+      EraRes ety eterm <- eraseElim ctx pair
+      let ty = sub1 (ret // shift 2) $
+               Ann (Pair (BVT 0 loc) (BVT 1 loc) loc) (weakT 2 ety) loc
+      (tfst, tsnd) <- wrapExpect `(expectSig) ctx loc ety
       let ctx' = extendTyN [< (qty, x, tfst), (qty, y, tsnd.term)] ctx
       body' <- eraseTerm ctx' ty body.term
-  let x' = ifErased qty (Erased loc) (Fst epair.term loc)
+      p <- mnb "p" loc
-      y' = ifErased qty (Erased loc) (Snd epair.term loc)
+      pure $ EraRes (sub1 ret pair) $
-  pure $ EraRes (sub1 ret pair) $ body' // fromSnocVect [< x', y']
+        Let p eterm
+          (Let x (Fst (B VZ loc) loc)
+            (Let y (Snd (B (VS VZ) loc) loc)
+              (body' // (B VZ loc ::: B (VS VZ) loc ::: shift 3))
+              loc) loc) loc
+    Erased => do
+      ety <- computeElimType ctx SOne pair
+      let ty = sub1 (ret // shift 2) $
+               Ann (Pair (BVT 0 loc) (BVT 1 loc) loc) (weakT 2 ety) loc
+      (tfst, tsnd) <- wrapExpect `(expectSig) ctx loc ety
+      let ctx' = extendTyN0 [< (x, tfst), (y, tsnd.term)] ctx
+      body' <- eraseTerm ctx' ty body.term
+      pure $ EraRes (sub1 ret pair) $ subN [< Erased loc, Erased loc] body'
 
 --  e ⤋ e' ⇒ (x : A) × B
 -- ----------------------
@@ -296,29 +306,30 @@ eraseElim ctx (Snd pair loc) = do
   b          <- snd <$> wrapExpect `(expectSig) ctx loc epair.type
   pure $ EraRes (sub1 b (Fst pair loc)) $ Snd epair.term loc
 
--- case0 e return z ⇒ R of {} ⤋ absurd ⇒ R[e/z]
+-- caseρ e return z ⇒ R of {} ⤋ absurd ⇒ R[e/z]
 --
 --               s ⤋ s' ⇐ R[𝐚∷{𝐚}/z]
 -- -----------------------------------------------
 --  case0 e return z ⇒ R of {𝐚 ⇒ s} ⤋ s' ⇒ R[e/z]
 --
---           e ⤋ e' ⇒ A   ρ ≠ 0    sᵢ ⤋ s'ᵢ ⇐ R[𝐚ᵢ/z]
+--         e ⤋ e' ⇒ A   sᵢ ⤋ s'ᵢ ⇐ R[𝐚ᵢ/z]   ρ ≠ 0   i ≠ 0
 -- -------------------------------------------------------------------
 --  caseρ e return z ⇒ R of {𝐚ᵢ ⇒ sᵢ} ⤋ case e of {𝐚ᵢ ⇒ s'ᵢ} ⇒ R[e/z]
-eraseElim ctx e@(CaseEnum qty tag ret arms loc) =
+eraseElim ctx e@(CaseEnum qty tag ret arms loc) = do
+  let ty = sub1 ret tag
   case isErased qty of
     Erased => case SortedMap.toList arms of
-      []         => pure $ EraRes (sub1 ret tag) $ Absurd loc
+      []         => pure $ EraRes ty $ Absurd loc
-      [(t, arm)] => do
+      [(t, rhs)] => do
-        let ty  = sub1 ret tag
+        let ty' = sub1 ret (Ann (Tag t loc) (enum [t] loc) loc)
-            ty' = sub1 ret (Ann (Tag t loc) (enum [t] loc) loc)
+        rhs' <- eraseTerm ctx ty' rhs
-        arm' <- eraseTerm ctx ty' arm
+        pure $ EraRes ty rhs'
-        pure $ EraRes ty arm'
       _ => throw $ CompileTimeOnly ctx $ E e
-    Kept => do
+    Kept => case List1.fromList $ SortedMap.toList arms of
-      let ty = sub1 ret tag
+      Nothing   => pure $ EraRes ty $ Absurd loc
+      Just arms => do
         etag <- eraseElim ctx tag
-      arms <- for (SortedMap.toList arms) $ \(t, rhs) => do
+        arms <- for arms $ \(t, rhs) => do
           let ty' = sub1 ret (Ann (Tag t loc) etag.type loc)
           rhs' <- eraseTerm ctx ty' rhs
           pure (t, rhs')
@@ -352,13 +363,12 @@ eraseElim ctx (CaseNat qty qtyIH nat ret zero succ loc) = do
                Erased => NSNonrec p    (sub1 (Erased loc) succ')
   pure $ EraRes ty $ CaseNat enat.term zero succ loc
 
---                b ⤋ b' ⇒ [π.A]   π ≠ 0
+--  b ⤋ b' ⇒ [π.A]   πρ ≠ 0   x : A | πρ.x ⊢ s ⤋ s' ⇐ R[[x]∷[π.A]/z]
---         x : A | πρ.x ⊢ s ⤋ s' ⇐ R[[x]∷[π.A]/z]
+-- ------------------------------------------------------------------
--- -------------------------------------------------------
+--   caseρ b return z ⇒ R of {[x] ⇒ s} ⤋ let x = b' in s' ⇒ R[b/z]
---  caseρ b return z ⇒ R of {[x] ⇒ s} ⤋ s'[b'/x] ⇒ R[b/z]
 --
---    b ⇒ [0.A]   x : A | 0.x ⊢ s ⤋ s' ⇐ R[[x]∷[0.A]/z]
+--   b ⇒ [π.A]   x : A | 0.x ⊢ s ⤋ s' ⇐ R[[x]∷[0.A]/z]   πρ = 0
--- -------------------------------------------------------
+-- -------------------------------------------------------------
 --  caseρ b return z ⇒ R of {[x] ⇒ s} ⤋ s'[⌷/x] ⇒ R[b/z]
 eraseElim ctx (CaseBox qty box ret body loc) = do
   tbox <- computeElimType ctx SOne box -- [fixme] is there any way to avoid this?
@@ -366,11 +376,11 @@ eraseElim ctx (CaseBox qty box ret body loc) = do
   let ctx' = extendTy (pi * qty) body.name tinner ctx
       bty  = sub1 (ret // shift 1) $
              Ann (Box (BVT 0 loc) loc) (weakT 1 tbox) loc
-  case isErased pi of
+  case isErased $ qty * pi of
     Kept => do
       ebox  <- eraseElim ctx box
       ebody <- eraseTerm ctx' bty body.term
-      pure $ EraRes (sub1 ret box) $ ebody // one ebox.term
+      pure $ EraRes (sub1 ret box) $ Let body.name ebox.term ebody loc
     Erased => do
       body' <- eraseTerm ctx' bty body.term
       pure $ EraRes (sub1 ret box) $ body' // one (Erased loc)
@@ -412,6 +422,67 @@ eraseElim ctx (DCloE (Sub term th)) =
   eraseElim ctx $ pushSubstsWith' th id term
 
 
+export
+uses : Var n -> Term n -> Nat
+uses i (F x _) = 0
+uses i (B j _) = if i == j then 1 else 0
+uses i (Lam x body _) = uses (VS i) body
+uses i (App fun arg _) = uses i fun + uses i arg
+uses i (Pair fst snd _) = uses i fst + uses i snd
+uses i (Fst pair _) = uses i pair
+uses i (Snd pair _) = uses i pair
+uses i (Tag tag _) = 0
+uses i (CaseEnum tag cases _) =
+  uses i tag + foldl max 0 (map (assert_total uses i . snd) cases)
+uses i (Absurd _) = 0
+uses i (Zero _) = 0
+uses i (Succ nat _) = uses i nat
+uses i (CaseNat nat zer suc _) = uses i nat + max (uses i zer) (uses' suc)
+  where uses' : CaseNatSuc n -> Nat
+        uses' (NSRec  _ _ s) = uses (VS (VS i)) s
+        uses' (NSNonrec _ s) = uses (VS i) s
+uses i (Let x rhs body _) = uses i rhs + uses (VS i) body
+uses i (Erased _) = 0
+
+export
+inlineable : U.Term n -> Bool
+inlineable (F {})      = True
+inlineable (B {})      = True
+inlineable (Tag {})    = True
+inlineable (Absurd {}) = True
+inlineable (Erased {}) = True
+inlineable _           = False
+
+export
+trimLets : U.Term n -> U.Term n
+trimLets (F x loc) = F x loc
+trimLets (B i loc) = B i loc
+trimLets (Lam x body loc) = Lam x (trimLets body) loc
+trimLets (App fun arg loc) = App (trimLets fun) (trimLets arg) loc
+trimLets (Pair fst snd loc) = Pair (trimLets fst) (trimLets snd) loc
+trimLets (Fst pair loc) = Fst (trimLets pair) loc
+trimLets (Snd pair loc) = Snd (trimLets pair) loc
+trimLets (Tag tag loc) = Tag tag loc
+trimLets (CaseEnum tag cases loc) =
+  CaseEnum (trimLets tag)
+           (map (map $ \c => trimLets $ assert_smaller cases c) cases) loc
+trimLets (Absurd loc) = Absurd loc
+trimLets (Zero loc) = Zero loc
+trimLets (Succ nat loc) = Succ (trimLets nat) loc
+trimLets (CaseNat nat zer suc loc) =
+  CaseNat (trimLets nat) (trimLets zer) (trimLets' suc) loc
+  where trimLets' : CaseNatSuc n -> CaseNatSuc n
+        trimLets' (NSRec x ih s) = NSRec    x ih $ trimLets s
+        trimLets' (NSNonrec x s) = NSNonrec x    $ trimLets s
+trimLets (Let x rhs body loc) =
+  let rhs'  = trimLets rhs
+      body' = trimLets body in
+  if inlineable rhs' || uses VZ body' == 1
+     then sub1 rhs' body'
+     else Let x rhs' body' loc
+trimLets (Erased loc) = Erased loc
+
+
 export covering
 eraseDef : Name -> Q.Definition -> Eff Erase U.Definition
 eraseDef name def@(MkDef qty type body loc) =
@@ -420,4 +491,4 @@ eraseDef name def@(MkDef qty type body loc) =
     Erased => pure ErasedDef
     Kept   => case body of
       Postulate     => throw $ Postulate loc name
-      Concrete body => KeptDef <$> eraseTerm empty type body
+      Concrete body => KeptDef . trimLets <$> eraseTerm empty type body

lib/Quox/Untyped/Syntax.idr

@@ -35,8 +35,8 @@ data Term where
   Snd : (pair : Term n) -> Loc -> Term n
 
   Tag : (tag : String) -> Loc -> Term n
-  CaseEnum : (tag : Term n) -> (cases : List (String, Term n)) -> Loc -> Term n
+  CaseEnum : (tag : Term n) -> (cases : List1 (String, Term n)) -> Loc -> Term n
-  ||| empty match with an erased head
+  ||| empty match
   Absurd : Loc -> Term n
 
   Zero : Loc -> Term n
@@ -47,6 +47,9 @@ data Term where
             Loc ->
             Term n
 
+  Let : (x : BindName) -> (rhs : Term n) -> (body : Term (S n)) -> Loc ->
+        Term n
+
   Erased : Loc -> Term n
 %name Term s, t, u
 
@@ -61,19 +64,20 @@ data CaseNatSuc where
 
 export
 Located (Term n) where
-  (F x loc).loc                 = loc
+  (F _ loc).loc           = loc
-  (B i loc).loc                 = loc
+  (B _ loc).loc           = loc
-  (Lam x body loc).loc          = loc
+  (Lam _ _ loc).loc       = loc
-  (App fun arg loc).loc         = loc
+  (App _ _ loc).loc       = loc
-  (Pair fst snd loc).loc        = loc
+  (Pair _ _ loc).loc      = loc
-  (Fst pair loc).loc            = loc
+  (Fst _ loc).loc         = loc
-  (Snd pair loc).loc            = loc
+  (Snd _ loc).loc         = loc
-  (Tag tag loc).loc             = loc
+  (Tag _ loc).loc         = loc
-  (CaseEnum tag cases loc).loc  = loc
+  (CaseEnum _ _ loc).loc  = loc
   (Absurd loc).loc        = loc
   (Zero loc).loc          = loc
-  (Succ nat loc).loc            = loc
+  (Succ _ loc).loc        = loc
-  (CaseNat nat zer suc loc).loc = loc
+  (CaseNat _ _ _ loc).loc = loc
+  (Let _ _ _ loc).loc     = loc
   (Erased loc).loc        = loc
 
 
@@ -85,6 +89,11 @@ public export
 Definitions = SortedMap Name Definition
 
 
+export
+letD, inD : {opts : LayoutOpts} -> Eff Pretty (Doc opts)
+letD = hl Syntax "let"
+inD  = hl Syntax "in"
+
 export
 prettyTerm : {opts : LayoutOpts} -> BContext n ->
              Term n -> Eff Pretty (Doc opts)
@@ -136,6 +145,33 @@ private
 sucPat : {opts : LayoutOpts} -> BindName -> Eff Pretty (Doc opts)
 sucPat x = pure $ !succD <++> !(prettyTBind x)
 
+private
+splitLam : Telescope' BindName a b -> Term b ->
+           Exists $ \c => (Telescope' BindName a c, Term c)
+splitLam ys (Lam x body _) = splitLam (ys :< x) body
+splitLam ys t              = Evidence _ (ys, t)
+
+private
+splitLet : Telescope (\i => (BindName, Term i)) a b -> Term b ->
+           Exists $ \c => (Telescope (\i => (BindName, Term i)) a c, Term c)
+splitLet ys (Let x rhs body _) = splitLet (ys :< (x, rhs)) body
+splitLet ys t                  = Evidence _ (ys, t)
+
+private
+prettyLets : {opts : LayoutOpts} ->
+             BContext a -> Telescope (\i => (BindName, Term i)) a b ->
+             Eff Pretty (SnocList (Doc opts))
+prettyLets xs lets = sequence $ snd $ go lets where
+  go : forall b. Telescope (\i => (BindName, Term i)) a b ->
+       (BContext b, SnocList (Eff Pretty (Doc opts)))
+  go [<]                = (xs, [<])
+  go (lets :< (x, rhs)) =
+    let (ys, docs) = go lets
+        doc = hsep <$> sequence
+              [letD, prettyTBind x, cstD, assert_total prettyTerm ys rhs, inD]
+    in
+    (ys :< x, docs :< doc)
+
 private
 sucCaseArm : {opts : LayoutOpts} ->
              CaseNatSuc n -> Eff Pretty (PrettyCaseArm (Doc opts) n)
@@ -148,9 +184,10 @@ prettyTerm _  (F x _) = prettyFree x
 prettyTerm xs (B i _) = prettyTBind $ xs !!! i
 prettyTerm xs (Lam x body _) =
   parensIfM Outer =<< do
-    header <- hsep <$> sequence [lamD, prettyTBind x, darrowD]
+    let Evidence n' (ys, body) = splitLam [< x] body
-    body   <- withPrec Outer $ prettyTerm (xs :< x) body
+    vars <- hsep . toList' <$> traverse prettyTBind ys
-    hangDSingle header body
+    body <- withPrec Outer $ assert_total prettyTerm (xs . ys) body
+    hangDSingle (hsep [!lamD, vars, !darrowD]) body
 prettyTerm xs (App fun arg _) = prettyApp xs fun arg
 prettyTerm xs (Pair fst snd _) =
   parens =<< separateTight !commaD <$>
@@ -161,21 +198,29 @@ prettyTerm xs (Tag tag _) = prettyTag tag
 prettyTerm xs (CaseEnum tag cases _) =
   assert_total
   prettyCase xs prettyTag tag $
-    map (\(t, rhs) => MkPrettyCaseArm t [] rhs) cases
+    map (\(t, rhs) => MkPrettyCaseArm t [] rhs) $ toList cases
 prettyTerm xs (Absurd _) = hl Syntax "absurd"
 prettyTerm xs (Zero _) = zeroD
 prettyTerm xs (Succ nat _) = prettyApp' xs !succD nat
 prettyTerm xs (CaseNat nat zer suc _) =
   assert_total
   prettyCase xs pure nat [MkPrettyCaseArm !zeroD [] zer, !(sucCaseArm suc)]
+prettyTerm xs (Let x rhs body _) =
+  parensIfM Outer =<< do
+    let Evidence n' (lets, body) = splitLet [< (x, rhs)] body
+    heads <- prettyLets xs lets
+    body  <- withPrec Outer $ assert_total prettyTerm (xs . map fst lets) body
+    let lines = toList $ heads :< body
+    pure $ ifMultiline (hsep lines) (vsep lines)
 prettyTerm _ (Erased _) =
   hl Syntax =<< ifUnicode "⌷" "[]"
 
 export
 prettyDef : {opts : LayoutOpts} -> Name ->
-            Definition -> Eff Pretty (Maybe (Doc opts))
+            Definition -> Eff Pretty (Doc opts)
-prettyDef _    ErasedDef     = [|Nothing|]
+prettyDef name ErasedDef =
-prettyDef name (KeptDef rhs) = map Just $ do
+  pure $ hsep [!(prettyFree name), !cstD, !(prettyTerm [<] $ Erased noLoc)]
+prettyDef name (KeptDef rhs) = do
   name <- prettyFree name
   eq   <- cstD
   rhs  <- prettyTerm [<] rhs
@@ -220,6 +265,8 @@ CanSubstSelf Term where
     CaseNat nat zer suc loc =>
       CaseNat (nat // th) (zer // th)
               (assert_total substSuc suc th) loc
+    Let x rhs body loc =>
+      Let x (rhs // th) (assert_total $ body // push th) loc
     Erased loc =>
       Erased loc
   where