From 15f6f4c8a42fca6886cad65c09186b4ca73fd11b Mon Sep 17 00:00:00 2001
From: rhiannon morris <rhi@rhiannon.website>
Date: Sat, 1 Apr 2023 16:01:53 +0200
Subject: [PATCH] fix nat elim quantities

---
 lib/Quox/Syntax/Qty.idr       | 20 ++++++++++++++++----
 lib/Quox/Syntax/Qty/Three.idr |  6 ++++++
 lib/Quox/Typechecker.idr      |  7 ++-----
 tests/Tests/Typechecker.idr   | 17 +++++++++++++++++
 4 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/lib/Quox/Syntax/Qty.idr b/lib/Quox/Syntax/Qty.idr
index 4e597c0..1612e75 100644
--- a/lib/Quox/Syntax/Qty.idr
+++ b/lib/Quox/Syntax/Qty.idr
@@ -10,9 +10,10 @@ import Data.DPair
 
 public export
 interface Eq q => IsQty q where
-  zero, one : q
-  (+),  (*) : q -> q -> q
-  lub       : q -> q -> Maybe q
+  zero, one, any : q
+
+  (+), (*) : q -> q -> q
+  lub      : q -> q -> Maybe q
 
   ||| true if bindings of this quantity will be erased
   ||| and must not be runtime relevant
@@ -20,6 +21,12 @@ interface Eq q => IsQty q where
   isZero : Dec1 IsZero
   zeroIsZero : IsZero zero
 
+  ||| true if bindings of this quantity can be used any number of times.
+  ||| this is needed for natural elimination
+  IsAny : Pred q
+  isAny : Dec1 IsAny
+  anyIsAny : IsAny any
+
   ||| ``p `Compat` q`` if it is ok for a binding of quantity `q` to be used
   ||| exactly `p` times. e.g. ``1 `Compat` 1``, ``1 `Compat` ω``.
   ||| if ``π `lub` ρ`` exists, then both `π` and `ρ` must be compatible with it
@@ -38,7 +45,8 @@ interface Eq q => IsQty q where
   ||| quantity. so not exact usage counts, maybe.
   IsGlobal : Pred q
   isGlobal : Dec1 IsGlobal
-  zeroIsGlobal : forall pi. IsZero pi -> IsGlobal zero
+  zeroIsGlobal : forall pi. IsZero pi -> IsGlobal pi
+  anyIsGlobal : forall pi. IsAny pi -> IsGlobal pi
 
   ||| prints in a form that can be a suffix of "case"
   prettySuffix : Pretty.HasEnv m => q -> m (Doc HL)
@@ -74,6 +82,10 @@ public export %inline
 gzero : IsQty q => GQty q
 gzero = Element zero $ zeroIsGlobal zeroIsZero
 
+public export %inline
+gany : IsQty q => GQty q
+gany = Element any $ anyIsGlobal anyIsAny
+
 export %inline
 globalToSubj : IsQty q => GQty q -> SQty q
 globalToSubj q = if isYes $ isZero q.fst then szero else sone
diff --git a/lib/Quox/Syntax/Qty/Three.idr b/lib/Quox/Syntax/Qty/Three.idr
index 2c05a8f..2a7d64a 100644
--- a/lib/Quox/Syntax/Qty/Three.idr
+++ b/lib/Quox/Syntax/Qty/Three.idr
@@ -106,6 +106,7 @@ public export
 IsQty Three where
   zero = Zero
   one  = One
+  any  = Any
 
   (+) = plus
   (*) = times
@@ -116,6 +117,10 @@ IsQty Three where
   isZero     = decEq Zero
   zeroIsZero = Refl
 
+  IsAny    = Equal Three.Any
+  isAny    = decEq Any
+  anyIsAny = Refl
+
   Compat = Compat3
   compat = compat3
 
@@ -128,6 +133,7 @@ IsQty Three where
   IsGlobal     = IsGlobal3
   isGlobal     = isGlobal3
   zeroIsGlobal = \Refl => GZero
+  anyIsGlobal  = \Refl => GAny
 
   prettySuffix = pretty0M
 
diff --git a/lib/Quox/Typechecker.idr b/lib/Quox/Typechecker.idr
index 824a442..f3a966a 100644
--- a/lib/Quox/Typechecker.idr
+++ b/lib/Quox/Typechecker.idr
@@ -386,16 +386,13 @@ parameters {auto _ : IsQty q}
         sucCtx    = extendTyN [< (pisg, p, Nat), (pi', ih, ret.term)] ctx
         sucType   = substCaseSuccRet ret
     sucout :< qp :< qih <- checkC sucCtx sg suc.term sucType
-    let Just armout = lubs ctx [zerout, sucout]
-          | _ => throw $ BadCaseQtys ctx $
-              [(zerout, Zero), (sucout, Succ $ FT $ unq p)]
     expectCompatQ qih (pi' * sg.fst)
     -- [fixme] better error here
     expectCompatQ (qp + qih) pisg
-    -- then Ψ | Γ ⊢ case ⋯ ⇒ A[n] ⊳ πΣn + (Σz ∧ Σs)
+    -- then Ψ | Γ ⊢ case ⋯ ⇒ A[n] ⊳ πΣn + Σz + ωΣs
     pure $ InfRes {
       type = sub1 ret n,
-      qout = pi * nres.qout + armout
+      qout = pi * nres.qout + zerout + any * sucout
     }
 
   infer' ctx sg (CaseBox pi box ret body) = do
diff --git a/tests/Tests/Typechecker.idr b/tests/Tests/Typechecker.idr
index deb04b7..202287a 100644
--- a/tests/Tests/Typechecker.idr
+++ b/tests/Tests/Typechecker.idr
@@ -401,6 +401,23 @@ tests = "typechecker" :- [
     todo "nat elim"
   ],
 
+  "natural elim" :- [
+    note "1 · λ n ⇒ case1 n return ℕ of { zero ⇒ 0; succ n ⇒ n }",
+    note "    ⇐ 1.ℕ → ℕ",
+    testTC "pred" $
+      check_ empty sone
+        ([< "n"] :\\ E (CaseNat One Zero (BV 0) (SN Nat)
+            Zero (SY [< "n", Unused] $ BVT 1)))
+        (Arr One Nat Nat),
+    note "1 · λ m n ⇒ case1 m return ℕ of { zero ⇒ n; succ _, 1.p ⇒ succ p }",
+    note "    ⇐ 1.ℕ → 1.ℕ → 1.ℕ",
+    testTC "plus" $
+      check_ empty sone
+        ([< "m", "n"] :\\ E (CaseNat One One (BV 1) (SN Nat)
+          (BVT 0) (SY [< Unused, "p"] $ Succ $ BVT 0)))
+        (Arr One Nat $ Arr One Nat Nat)
+  ],
+
   "box types" :- [
     testTC "0 · [0.ℕ] ⇐ ★₀" $
       check_ empty szero (BOX Zero Nat) (TYPE 0),