make nat elimination with erased IH non-recursive at runtime

lib/Quox/Untyped/Erase.idr

@@ -324,24 +324,33 @@ eraseElim ctx e@(CaseEnum qty tag ret arms loc) =
         pure (t, rhs')
       pure $ EraRes ty $ CaseEnum etag.term arms loc
 
---             n ⤋ n' ⇒ ℕ   z ⤋ z' ⇐ R[zero∷ℕ/z]
+--        n ⤋ n' ⇒ ℕ   z ⤋ z' ⇐ R[zero∷ℕ/z]   ς ≠ 0
 --  m : ℕ, ih : R[m/z] | ρ.m, ς.ih ⊢ s ⤋ s' ⇐ R[(succ m)∷ℕ/z]
 -- -----------------------------------------------------------
 --      caseρ n return z ⇒ R of {0 ⇒ z; succ m, ς.ih ⇒ s}
 --        ⤋
---      case n' of {0 ⇒ z'; succ m, ih ⇒ s'}
+--      case n' of {0 ⇒ z'; succ m, ih ⇒ s'} ⇒ R[n/z]
---        ⇒ R[n/z]
+--
+--             n ⤋ n' ⇒ ℕ   z ⤋ z' ⇐ R[zero∷ℕ/z]
+--  m : ℕ, ih : R[m/z] | ρ.m, 0.ih ⊢ s ⤋ s' ⇐ R[(succ m)∷ℕ/z]
+-- -----------------------------------------------------------
+--      caseρ n return z ⇒ R of {0 ⇒ z; succ m, 0.ih ⇒ s}
+--        ⤋
+--      case n' of {0 ⇒ z'; succ m ⇒ s'[⌷/ih]} ⇒ R[n/z]
 eraseElim ctx (CaseNat qty qtyIH nat ret zero succ loc) = do
   let ty = sub1 ret nat
   enat <- eraseElim ctx nat
   zero <- eraseTerm ctx (sub1 ret (Ann (Zero loc) (Nat loc) loc)) zero
   let [< p, ih] = succ.names
-  succ <- eraseTerm
+  succ' <- eraseTerm
     (extendTyN [< (qty, p, Nat loc),
                   (qtyIH, ih, sub1 (ret // shift 1) (BV 0 loc))] ctx)
     (sub1 (ret // shift 2) (Ann (Succ (BVT 1 loc) loc) (Nat loc) loc))
     succ.term
-  pure $ EraRes ty $ CaseNat enat.term zero p ih succ loc
+  let succ = case isErased qtyIH of
+               Kept   => NSRec    p ih succ'
+               Erased => NSNonrec p    (sub1 (Erased loc) succ')
+  pure $ EraRes ty $ CaseNat enat.term zero succ loc
 
 --                b ⤋ b' ⇒ [π.A]   π ≠ 0
 --         x : A | πρ.x ⊢ s ⤋ s' ⇐ R[[x]∷[π.A]/z]

lib/Quox/Untyped/Syntax.idr

@@ -18,7 +18,12 @@ import Derive.Prelude
 
 
 public export
-data Term : Nat -> Type where
+data Term : Nat -> Type
+
+public export
+data CaseNatSuc : Nat -> Type
+
+data Term where
   F : (x : Name) -> Loc -> Term n
   B : (i : Var n) -> Loc -> Term n
 
@@ -38,31 +43,38 @@ data Term : Nat -> Type where
   Succ : (nat : Term n) -> Loc -> Term n
   CaseNat : (nat : Term n) ->
             (zer : Term n) ->
-            (x, ih : BindName) -> (suc : Term (2 + n)) ->
+            (suc : CaseNatSuc n) ->
             Loc ->
             Term n
 
   Erased : Loc -> Term n
 %name Term s, t, u
-%runElab deriveIndexed "Term" [Eq, Ord, Show]
+
+data CaseNatSuc where
+  NSRec    : (x, ih : BindName) -> Term (2 + n) -> CaseNatSuc n
+  NSNonrec : (x     : BindName) -> Term (S n)   -> CaseNatSuc n
+%name CaseNatSuc suc
+
+%runElab deriveParam $
+  map (\ty => PI ty allIndices [Eq, Ord, Show]) ["Term", "CaseNatSuc"]
 
 
 export
 Located (Term n) where
-  (F x loc).loc                      = loc
+  (F x loc).loc                 = loc
-  (B i loc).loc                      = loc
+  (B i loc).loc                 = loc
-  (Lam x body loc).loc               = loc
+  (Lam x body loc).loc          = loc
-  (App fun arg loc).loc              = loc
+  (App fun arg loc).loc         = loc
-  (Pair fst snd loc).loc             = loc
+  (Pair fst snd loc).loc        = loc
-  (Fst pair loc).loc                 = loc
+  (Fst pair loc).loc            = loc
-  (Snd pair loc).loc                 = loc
+  (Snd pair loc).loc            = loc
-  (Tag tag loc).loc                  = loc
+  (Tag tag loc).loc             = loc
-  (CaseEnum tag cases loc).loc       = loc
+  (CaseEnum tag cases loc).loc  = loc
-  (Absurd loc).loc                   = loc
+  (Absurd loc).loc              = loc
-  (Zero loc).loc                     = loc
+  (Zero loc).loc                = loc
-  (Succ nat loc).loc                 = loc
+  (Succ nat loc).loc            = loc
-  (CaseNat nat zer x ih suc loc).loc = loc
+  (CaseNat nat zer suc loc).loc = loc
-  (Erased loc).loc                   = loc
+  (Erased loc).loc              = loc
 
 
 public export
@@ -120,6 +132,18 @@ prettyCase xs f head arms =
       (hsep [header, lb, separateTight sc cases, rb])
       (vsep [hsep [header, lb], indent d $ vsep (map (<+> sc) cases), rb])
 
+private
+sucPat : {opts : LayoutOpts} -> BindName -> Eff Pretty (Doc opts)
+sucPat x = pure $ !succD <++> !(prettyTBind x)
+
+private
+sucCaseArm : {opts : LayoutOpts} ->
+             CaseNatSuc n -> Eff Pretty (PrettyCaseArm (Doc opts) n)
+sucCaseArm (NSRec x ih s) = pure $
+  MkPrettyCaseArm (!(sucPat x) <+> !commaD <++> !(prettyTBind ih)) [x, ih] s
+sucCaseArm (NSNonrec x s) = pure $
+  MkPrettyCaseArm !(sucPat x) [x] s
+
 prettyTerm _  (F x _) = prettyFree x
 prettyTerm xs (B i _) = prettyTBind $ xs !!! i
 prettyTerm xs (Lam x body _) =
@@ -141,14 +165,9 @@ prettyTerm xs (CaseEnum tag cases _) =
 prettyTerm xs (Absurd _) = hl Syntax "absurd"
 prettyTerm xs (Zero _) = zeroD
 prettyTerm xs (Succ nat _) = prettyApp' xs !succD nat
-prettyTerm xs (CaseNat nat zer x ih suc _) =
+prettyTerm xs (CaseNat nat zer suc _) =
   assert_total
-  prettyCase xs pure nat
+  prettyCase xs pure nat [MkPrettyCaseArm !zeroD [] zer, !(sucCaseArm suc)]
-    [MkPrettyCaseArm !zeroD  []      zer,
-     MkPrettyCaseArm !sucPat [x, ih] suc]
-where
-  sucPat = pure $
-    hsep [!succD, !(prettyTBind x) <+> !commaD, !(prettyTBind ih)]
 prettyTerm _ (Erased _) =
   hl Syntax =<< ifUnicode "⌷" "[]"
 
@@ -198,11 +217,16 @@ CanSubstSelf Term where
       Zero loc
     Succ nat loc =>
       Succ (nat // th) loc
-    CaseNat nat zer x ih suc loc =>
+    CaseNat nat zer suc loc =>
       CaseNat (nat // th) (zer // th)
-              x ih (assert_total $ suc // pushN 2 th) loc
+              (assert_total substSuc suc th) loc
     Erased loc =>
       Erased loc
+  where
+    substSuc : forall from, to.
+               CaseNatSuc from -> USubst from to -> CaseNatSuc to
+    substSuc (NSRec    x ih t) th = NSRec    x ih $ t // pushN 2 th
+    substSuc (NSNonrec x    t) th = NSNonrec x    $ t // push    th
 
 public export
 subN : SnocVect s (Term n) -> Term (s + n) -> Term n