From 7f873ceec89756f9f8b898dccbadc17abbed8433 Mon Sep 17 00:00:00 2001
From: flupe <lucas@escot.me>
Date: Wed, 7 Aug 2024 14:45:15 +0200
Subject: [PATCH] further work on compiling unboxed datatypes

---
 lib/Haskell/Extra/Singleton.agda |  4 ++--
 src/Agda2Hs/Compile.hs           |  1 +
 src/Agda2Hs/Compile/Function.hs  | 11 ++++++++---
 src/Agda2Hs/Compile/Record.hs    |  7 ++++---
 src/Agda2Hs/Compile/Type.hs      | 29 +++++++++++++++++++++++++----
 src/Agda2Hs/Compile/Utils.hs     | 19 +++++++++++++++++--
 test/Singleton.agda              | 13 +++++++++++++
 7 files changed, 70 insertions(+), 14 deletions(-)
 create mode 100644 test/Singleton.agda

diff --git a/lib/Haskell/Extra/Singleton.agda b/lib/Haskell/Extra/Singleton.agda
index 1604e5c5..94300930 100644
--- a/lib/Haskell/Extra/Singleton.agda
+++ b/lib/Haskell/Extra/Singleton.agda
@@ -8,9 +8,9 @@ data Singleton (a : Set) : @0 a → Set where
 {-# COMPILE AGDA2HS Singleton unboxed #-}
 
 pureSingl : {a : Set} (x : a) → Singleton a x
-pureSingl = MkSingl
+pureSingl x = MkSingl x
 
-{-# COMPILE AGDA2HS pureSingl inline #-}
+{-# COMPILE AGDA2HS pureSingl transparent #-}
 
 fmapSingl
   : {a b : Set} (f : a → b) {@0 x : a}
diff --git a/src/Agda2Hs/Compile.hs b/src/Agda2Hs/Compile.hs
index 22c93619..b3172924 100644
--- a/src/Agda2Hs/Compile.hs
+++ b/src/Agda2Hs/Compile.hs
@@ -89,6 +89,7 @@ compile opts tlm _ def =
         (NoPragma            , _         ) -> return []
         (ExistingClassPragma , _         ) -> return []
         (UnboxPragma s       , Record{}  ) -> [] <$ checkUnboxPragma def
+        (UnboxPragma s       , Datatype{}) -> [] <$ checkUnboxPragma def
         (TransparentPragma   , Function{}) -> [] <$ checkTransparentPragma def
         (InlinePragma        , Function{}) -> [] <$ checkInlinePragma def
         (TuplePragma b       , Record{}  ) -> return []
diff --git a/src/Agda2Hs/Compile/Function.hs b/src/Agda2Hs/Compile/Function.hs
index 0fb65335..465095b9 100644
--- a/src/Agda2Hs/Compile/Function.hs
+++ b/src/Agda2Hs/Compile/Function.hs
@@ -399,7 +399,7 @@ checkInlinePragma def@Defn{defName = f} = do
     [c] ->
       unlessM (allowedPats (namedClausePats c)) $ genericDocError =<<
         "Cannot make function" <+> prettyTCM (defName def) <+> "inlinable." <+>
-        "Inline functions can only use variable patterns or transparent record constructor patterns."
+        "Inline functions can only use variable patterns or unboxed constructor patterns."
     _ ->
       genericDocError =<<
         "Cannot make function" <+> prettyTCM f <+> "inlinable." <+>
@@ -408,11 +408,16 @@ checkInlinePragma def@Defn{defName = f} = do
   where allowedPat :: DeBruijnPattern -> C Bool
         allowedPat VarP{} = pure True
         -- only allow matching on (unboxed) record constructors
-        allowedPat (ConP ch ci cargs) =
+        allowedPat (ConP ch ci cargs) = do
+          reportSDoc "agda2hs.compile.inline" 18 $ "Checking unboxing of constructor: "<+> prettyTCM ch
           isUnboxConstructor (conName ch) >>= \case
             Just _  -> allowedPats cargs
             Nothing -> pure False
-        allowedPat _ = pure False
+        -- NOTE(flupe): dot patterns should really only be allowed for *erased* arguments, I think.
+        allowedPat (DotP _ _) = pure True
+        allowedPat p = do
+          reportSDoc "agda2hs.compile.inline" 18 $ "Unsupported pattern"  <+> prettyTCM p
+          pure False
 
         allowedPats :: NAPs -> C Bool
         allowedPats pats = allM pats (allowedPat . dget . dget)
diff --git a/src/Agda2Hs/Compile/Record.hs b/src/Agda2Hs/Compile/Record.hs
index b82b0f4a..1045bc59 100644
--- a/src/Agda2Hs/Compile/Record.hs
+++ b/src/Agda2Hs/Compile/Record.hs
@@ -171,9 +171,7 @@ compileRecord target def = do
       return $ Hs.DataDecl () don Nothing hd [conDecl] ds
 
 checkUnboxPragma :: Definition -> C ()
-checkUnboxPragma def = do
-  let Record{..} = theDef def
-
+checkUnboxPragma def | Record{..} <- theDef def = do
   -- recRecursive can be used again after agda 2.6.4.2 is released
   -- see agda/agda#7042
   unless (all null recMutual) $ genericDocError
@@ -197,3 +195,6 @@ checkUnboxPragma def = do
       DOType -> genericDocError =<< text "Type field in unboxed record not supported"
       DOInstance -> genericDocError =<< text "Instance field in unboxed record not supported"
       DOTerm -> (absName tel:) <$> underAbstraction a tel nonErasedFields
+
+checkUnboxPragma def | Datatype{..} <- theDef def = do
+  pure () -- TOOD(flupe):
diff --git a/src/Agda2Hs/Compile/Type.hs b/src/Agda2Hs/Compile/Type.hs
index 12ac4405..2dd9d8ff 100644
--- a/src/Agda2Hs/Compile/Type.hs
+++ b/src/Agda2Hs/Compile/Type.hs
@@ -9,6 +9,7 @@ import Control.Monad.Trans ( lift )
 import Control.Monad.Reader ( asks )
 import Data.List ( find )
 import Data.Maybe ( mapMaybe, isJust )
+import Data.Foldable ( toList )
 import qualified Data.Set as Set ( singleton )
 
 import qualified Language.Haskell.Exts.Syntax as Hs
@@ -21,6 +22,7 @@ import Agda.Syntax.Common
 import Agda.Syntax.Internal
 import Agda.Syntax.Common.Pretty ( prettyShow )
 
+import Agda.TypeChecking.Datatypes
 import Agda.TypeChecking.Pretty
 import Agda.TypeChecking.Reduce ( reduce, unfoldDefinitionStep )
 import Agda.TypeChecking.Substitute
@@ -76,7 +78,7 @@ delayType [] = genericDocError =<< text "Cannot compile unapplied Delay type"
 compileTypeWithStrictness :: Term -> C (Strictness, Hs.Type ())
 compileTypeWithStrictness t = do
   s <- case t of
-    Def f es -> fromMaybe Lazy <$> isUnboxRecord f
+    Def f es -> fromMaybe Lazy <$> isUnboxType f
     _        -> return Lazy
   ty <- compileType t
   pure (s, ty)
@@ -106,7 +108,7 @@ compileType t = do
               <+> text "in Haskell type"
          | Just semantics <- isSpecialType f -> setCurrentRange f $ semantics es
          | Just args <- allApplyElims es ->
-             ifJustM (isUnboxRecord f) (\_ -> compileUnboxType f args) $
+             ifJustM (isUnboxType f) (\_ -> compileUnboxType f args) $
              ifJustM (isTupleRecord f) (\b -> compileTupleType f b args) $
              ifM (isTransparentFunction f) (compileTransparentType (defType def) args) $
              ifM (isInlinedFunction f) (compileInlineType f es) $ do
@@ -174,11 +176,31 @@ compileTelSize (ExtendTel a tel) = compileDom a >>= \case
 compileUnboxType :: QName -> Args -> C (Hs.Type ())
 compileUnboxType r pars = do
   def <- getConstInfo r
-  let tel = recTel (theDef def) `apply` pars
+
+  tel <- case theDef def of
+
+    Record{recTel} -> pure (recTel `apply` pars)
+
+    Datatype{dataCons = [con]} -> do
+      -- NOTE(flupe): con is the *only* constructor of the unboxed datatype
+      Constructor { conSrcCon } <- theDef <$> getConstInfo con
+      -- we retrieve its type
+      Just (_, ty) <- getConType conSrcCon (El (DummyS "some level") (apply (Def r []) pars))
+      theTel <$> telView ty
+
   compileTel tel >>= \case
     [t] -> return t
     _   -> __IMPOSSIBLE__
 
+  where
+    compileTel :: Telescope -> C [Hs.Type ()]
+    compileTel EmptyTel = return []
+    compileTel (ExtendTel a tel) = compileDom a >>= \case
+      DODropped  -> underAbstraction a tel compileTel
+      DOInstance -> __IMPOSSIBLE__
+      DOType     -> __IMPOSSIBLE__
+      DOTerm     -> (:) <$> compileType (unEl $ unDom a) <*> underAbstraction a tel compileTel
+
 compileTupleType :: QName -> Hs.Boxed -> Args -> C (Hs.Type ())
 compileTupleType r b pars = do
   tellUnboxedTuples b
@@ -192,7 +214,6 @@ compileTransparentType ty args = compileTypeArgs ty args >>= \case
   (v:vs) -> return $ v `tApp` vs
   []     -> __IMPOSSIBLE__
 
-
 compileInlineType :: QName -> Elims -> C (Hs.Type ())
 compileInlineType f args = do
   Function { funClauses = cs } <- theDef <$> getConstInfo f
diff --git a/src/Agda2Hs/Compile/Utils.hs b/src/Agda2Hs/Compile/Utils.hs
index a62f80a5..d1649325 100644
--- a/src/Agda2Hs/Compile/Utils.hs
+++ b/src/Agda2Hs/Compile/Utils.hs
@@ -192,9 +192,24 @@ isUnboxRecord q = do
         _             -> Nothing
     _ -> return Nothing
 
+isUnboxType :: QName -> C (Maybe Strictness)
+isUnboxType q = do
+  getConstInfo q >>= \case
+    Defn{defName = r} ->
+      processPragma r <&> \case
+        UnboxPragma s -> Just s
+        _             -> Nothing
+    -- _ -> return Nothing
+
 isUnboxConstructor :: QName -> C (Maybe Strictness)
-isUnboxConstructor q =
-  caseMaybeM (isRecordConstructor q) (return Nothing) $ isUnboxRecord . fst
+isUnboxConstructor q = do
+  def <- getConstInfo q
+  case theDef def of
+      Constructor {conData = r} -> isUnboxType r
+      _                         -> pure Nothing
+
+  -- caseMaybeM (isRecordConstructor) isUnboxType q
+  -- caseMaybeM (isRecordConstructor q) (return Nothing) $ isUnboxRecord . fst
 
 isUnboxProjection :: QName -> C (Maybe Strictness)
 isUnboxProjection q =
diff --git a/test/Singleton.agda b/test/Singleton.agda
new file mode 100644
index 00000000..3d60d5b4
--- /dev/null
+++ b/test/Singleton.agda
@@ -0,0 +1,13 @@
+open import Haskell.Prelude
+open import Haskell.Extra.Singleton
+
+test1 : (@0 x : ⊤) → Singleton ⊤ x
+test1 x = MkSingl tt
+
+{-# COMPILE AGDA2HS test1 #-}
+
+-- doesn't work because inline functions only reduce because of inline record transformation
+-- test2 : {a b : Set} (f : a → b) {@0 x : a} → Singleton a x → Singleton b (f x)
+-- test2 f sx = fmapSingl f sx
+-- 
+-- {-# COMPILE AGDA2HS test2 #-}