From 984821ab84b33917a1709c36bf66a0732ebbb46d Mon Sep 17 00:00:00 2001
From: Jesper Cockx <jesper@sikanda.be>
Date: Mon, 15 Jan 2024 18:31:04 +0100
Subject: [PATCH] Add built-in support for Delay monad

---
 lib/Haskell/Extra/Delay.agda      | 39 +++++++++++++++++++++++++++++++
 lib/Haskell/Extra/Loop.agda       |  8 -------
 lib/Haskell/Extra/Refinement.agda |  7 +++++-
 lib/Haskell/Prelude.agda          | 13 +++++++++--
 src/Agda2Hs/Compile/Function.hs   |  2 ++
 src/Agda2Hs/Compile/Term.hs       |  3 +++
 src/Agda2Hs/Compile/Type.hs       |  7 ++++++
 test/AllTests.agda                |  2 ++
 test/Delay.agda                   | 23 ++++++++++++++++++
 test/Fail/MatchOnDelay.agda       | 11 +++++++++
 test/golden/AllTests.hs           |  1 +
 test/golden/Delay.hs              |  7 ++++++
 test/golden/MatchOnDelay.err      |  2 ++
 13 files changed, 114 insertions(+), 11 deletions(-)
 create mode 100644 lib/Haskell/Extra/Delay.agda
 create mode 100644 test/Delay.agda
 create mode 100644 test/Fail/MatchOnDelay.agda
 create mode 100644 test/golden/Delay.hs
 create mode 100644 test/golden/MatchOnDelay.err

diff --git a/lib/Haskell/Extra/Delay.agda b/lib/Haskell/Extra/Delay.agda
new file mode 100644
index 00000000..40471b71
--- /dev/null
+++ b/lib/Haskell/Extra/Delay.agda
@@ -0,0 +1,39 @@
+{-# OPTIONS --sized-types #-}
+
+module Haskell.Extra.Delay where
+
+open import Agda.Builtin.Size public
+
+open import Haskell.Prelude
+open import Haskell.Prim.Thunk
+open import Haskell.Extra.Refinement
+
+private variable
+  x y z : a
+  @0 i : Size
+
+data Delay (a : Set) (@0 i : Size) : Set where
+  now : a → Delay a i
+  later : Thunk (Delay a) i → Delay a i
+
+data HasResult (x : a) : Delay a i → Set where
+  now   : HasResult x (now x)
+  later : HasResult x (y .force) → HasResult x (later y)
+
+runDelay : {@0 x : a} (y : Delay a ∞) → @0 HasResult x y → a
+runDelay (now x) now = x
+runDelay (later y) (later p) = runDelay (y .force) p
+
+runDelaySound : {@0 x : a} (y : Delay a ∞) → (@0 hr : HasResult x y) → runDelay y hr ≡ x
+runDelaySound (now x) now = refl
+runDelaySound (later y) (later hr) = runDelaySound (y .force) hr
+
+-- tryDelay and unDelay cannot and should not be compiled to Haskell,
+-- so they are marked as erased.
+@0 tryDelay : (y : Delay a ∞) → Nat → Maybe (∃ a (λ x → HasResult x y))
+tryDelay (now x)   _       = Just (x ⟨ now ⟩)
+tryDelay (later y) zero    = Nothing
+tryDelay (later y) (suc n) = fmap (mapRefine later) (tryDelay (y .force) n)
+
+@0 unDelay : (y : Delay a ∞) (n : Nat) → @0 {IsJust (tryDelay y n)} → a
+unDelay y n {p} = fromJust (tryDelay y n) {p} .value
diff --git a/lib/Haskell/Extra/Loop.agda b/lib/Haskell/Extra/Loop.agda
index cd908848..5009e80f 100644
--- a/lib/Haskell/Extra/Loop.agda
+++ b/lib/Haskell/Extra/Loop.agda
@@ -3,14 +3,6 @@ open import Haskell.Prelude
 
 module Haskell.Extra.Loop where
 
-IsJust : Maybe a → Set
-IsJust Nothing  = ⊥
-IsJust (Just _) = ⊤
-
-fromJust : (x : Maybe a) → @0 {IsJust x} → a
-fromJust (Just x) = x
-fromJust Nothing  = error "fromJust called on Nothing"
-
 data Fuel (f : a → Either a b) : (x : Either a b) → Set where
   done : ∀ {y} → Fuel f (Right y)
   more : ∀ {x} → Fuel f (f x) → Fuel f (Left x)
diff --git a/lib/Haskell/Extra/Refinement.agda b/lib/Haskell/Extra/Refinement.agda
index 52ee3e0a..f3211ee0 100644
--- a/lib/Haskell/Extra/Refinement.agda
+++ b/lib/Haskell/Extra/Refinement.agda
@@ -6,10 +6,15 @@ open import Agda.Primitive
 private variable
   ℓ ℓ′ : Level
 
-record ∃ (@0 a : Set ℓ) (@0 P : a → Set ℓ′) : Set (ℓ ⊔ ℓ′) where
+record ∃ (a : Set ℓ) (@0 P : a → Set ℓ′) : Set (ℓ ⊔ ℓ′) where
   constructor _⟨_⟩
   field
     value    : a
     @0 proof : P value
 open ∃ public
 {-# COMPILE AGDA2HS ∃ unboxed #-}
+
+mapRefine : {@0 P Q : a → Set ℓ} (@0 f : ∀ {x} → P x → Q x) → ∃ a P → ∃ a Q
+mapRefine f (x ⟨ p ⟩) = x ⟨ f p ⟩
+
+{-# COMPILE AGDA2HS mapRefine transparent #-}
diff --git a/lib/Haskell/Prelude.agda b/lib/Haskell/Prelude.agda
index 18d97843..dbedf71c 100644
--- a/lib/Haskell/Prelude.agda
+++ b/lib/Haskell/Prelude.agda
@@ -124,7 +124,16 @@ lookup : ⦃ Eq a ⦄ → a → List (a × b) → Maybe b
 lookup x []              = Nothing
 lookup x ((x₁ , y) ∷ xs) = if x == x₁ then Just y else lookup x xs
 
-private variable A B : Set
+-------------------------------------------------
+-- Unsafe functions
 
-coerce : @0 A ≡ B → A → B
+coerce : @0 a ≡ b → a → b
 coerce refl x = x
+
+IsJust : Maybe a → Set
+IsJust Nothing  = ⊥
+IsJust (Just _) = ⊤
+
+fromJust : (x : Maybe a) → @0 {IsJust x} → a
+fromJust Nothing  = error "fromJust Nothing"
+fromJust (Just x) = x
diff --git a/src/Agda2Hs/Compile/Function.hs b/src/Agda2Hs/Compile/Function.hs
index 40c1e83f..beeec77e 100644
--- a/src/Agda2Hs/Compile/Function.hs
+++ b/src/Agda2Hs/Compile/Function.hs
@@ -57,6 +57,8 @@ isSpecialPat qn = case prettyShow qn of
     badConstructors =
       [ "Agda.Builtin.Nat.Nat.zero"
       , "Agda.Builtin.Nat.Nat.suc"
+      , "Haskell.Extra.Delay.Delay.now"
+      , "Haskell.Extra.Delay.Delay.later"
       ]
 
 isUnboxCopattern :: DeBruijnPattern -> C Bool
diff --git a/src/Agda2Hs/Compile/Term.hs b/src/Agda2Hs/Compile/Term.hs
index 02478622..a3310b03 100644
--- a/src/Agda2Hs/Compile/Term.hs
+++ b/src/Agda2Hs/Compile/Term.hs
@@ -53,6 +53,7 @@ isSpecialTerm q = case prettyShow q of
   "Haskell.Prim.case_of_"                       -> Just caseOf
   "Haskell.Prim.Monad.Do.Monad._>>=_"           -> Just bind
   "Haskell.Prim.Monad.Do.Monad._>>_"            -> Just sequ
+  "Haskell.Extra.Delay.runDelay"                -> Just $ const compileErasedApp
   "Agda.Builtin.FromNat.Number.fromNat"         -> Just fromNat
   "Agda.Builtin.FromNeg.Negative.fromNeg"       -> Just fromNeg
   "Agda.Builtin.FromString.IsString.fromString" -> Just fromString
@@ -63,6 +64,8 @@ isSpecialCon = prettyShow >>> \case
   "Haskell.Prim.Tuple._,_"         -> Just tupleTerm
   "Haskell.Prim.Tuple._×_×_._,_,_" -> Just tupleTerm
   "Haskell.Extra.Erase.Erased"     -> Just (\_ _ _ -> erasedTerm)
+  "Haskell.Extra.Delay.Delay.now"   -> Just $ \_ _ -> compileErasedApp
+  "Haskell.Extra.Delay.Delay.later" -> Just $ \_ _ -> compileErasedApp
   _ -> Nothing
 
 tupleTerm :: ConHead -> ConInfo -> Elims -> C (Hs.Exp ())
diff --git a/src/Agda2Hs/Compile/Type.hs b/src/Agda2Hs/Compile/Type.hs
index 2c775b4e..ac2507c1 100644
--- a/src/Agda2Hs/Compile/Type.hs
+++ b/src/Agda2Hs/Compile/Type.hs
@@ -44,6 +44,7 @@ isSpecialType = prettyShow >>> \case
   "Haskell.Prim.Tuple._×_"   -> Just tupleType
   "Haskell.Prim.Tuple._×_×_" -> Just tupleType
   "Haskell.Extra.Erase.Erase" -> Just erasedType
+  "Haskell.Extra.Delay.Delay" -> Just delayType
   _ -> Nothing
 
 tupleType :: QName -> Elims -> C (Hs.Type ())
@@ -55,6 +56,12 @@ tupleType q es = do
 erasedType :: QName -> Elims -> C (Hs.Type ())
 erasedType _ _ = return $ Hs.TyTuple () Hs.Boxed []
 
+delayType :: QName -> Elims -> C (Hs.Type ())
+delayType _ (Apply a : _) = compileType (unArg a)
+delayType _ (_ : _) = __IMPOSSIBLE__
+delayType _ [] = genericDocError =<< text "Cannot compile unapplied Delay type"
+
+
 -- | Add a class constraint to a Haskell type.
 constrainType
   :: Hs.Asst () -- ^ The class assertion.
diff --git a/test/AllTests.agda b/test/AllTests.agda
index a606ff58..b4328359 100644
--- a/test/AllTests.agda
+++ b/test/AllTests.agda
@@ -67,6 +67,7 @@ import Coerce
 import Inlining
 import EraseType
 import Issue257
+import Delay
 
 {-# FOREIGN AGDA2HS
 import Issue14
@@ -133,4 +134,5 @@ import ModuleParametersImports
 import Coerce
 import Inlining
 import EraseType
+import Delay
 #-}
diff --git a/test/Delay.agda b/test/Delay.agda
new file mode 100644
index 00000000..b74ec943
--- /dev/null
+++ b/test/Delay.agda
@@ -0,0 +1,23 @@
+
+module Delay where
+
+open import Haskell.Prelude
+open import Haskell.Prim.Thunk
+open import Haskell.Extra.Delay
+
+open import Agda.Builtin.Size
+
+postulate
+  div : Int → Int → Int
+  mod : Int → Int → Int
+
+even : Int → Bool
+even x = mod x 2 == 0
+
+collatz : ∀ {@0 j} → Int → Delay Int j
+collatz x = 
+  if x == 0 then now 0
+  else if even x then later (λ where .force → collatz (div x 2))
+  else later λ where .force → collatz (3 * x + 1)
+
+{-# COMPILE AGDA2HS collatz #-}
diff --git a/test/Fail/MatchOnDelay.agda b/test/Fail/MatchOnDelay.agda
new file mode 100644
index 00000000..10ce5a12
--- /dev/null
+++ b/test/Fail/MatchOnDelay.agda
@@ -0,0 +1,11 @@
+
+module Fail.MatchOnDelay where
+
+open import Haskell.Prelude
+open import Haskell.Extra.Delay
+
+bad : Delay a ∞ → Bool
+bad (now x) = True
+bad (later x) = False
+
+{-# COMPILE AGDA2HS bad #-}
diff --git a/test/golden/AllTests.hs b/test/golden/AllTests.hs
index d69d9254..d38a514a 100644
--- a/test/golden/AllTests.hs
+++ b/test/golden/AllTests.hs
@@ -64,4 +64,5 @@ import ModuleParametersImports
 import Coerce
 import Inlining
 import EraseType
+import Delay
 
diff --git a/test/golden/Delay.hs b/test/golden/Delay.hs
new file mode 100644
index 00000000..6eeaad90
--- /dev/null
+++ b/test/golden/Delay.hs
@@ -0,0 +1,7 @@
+module Delay where
+
+collatz :: Int -> Int
+collatz x
+  = if x == 0 then 0 else
+      if even x then collatz (div x 2) else collatz (3 * x + 1)
+
diff --git a/test/golden/MatchOnDelay.err b/test/golden/MatchOnDelay.err
new file mode 100644
index 00000000..468a5242
--- /dev/null
+++ b/test/golden/MatchOnDelay.err
@@ -0,0 +1,2 @@
+test/Fail/MatchOnDelay.agda:7,1-4
+constructor `now` not supported in patterns