Merge pull request #114 from Plutonomicon/mario/records-on-staging

Mario/records on staging
Plutonomicon · Jan 14, 2022 · 231267e · 231267e
2 parents 7f66dee + 326b32f
commit 231267e
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 44 deletions.
diff --git a/Plutarch/Rec.hs b/Plutarch/Rec.hs
@@ -12,7 +12,7 @@ import Data.Functor.Compose (Compose)
 import Data.Kind (Type)
 import Data.Monoid (Dual (Dual, getDual), Endo (Endo, appEndo), Sum (Sum, getSum))
 import Numeric.Natural (Natural)
-import Plutarch (PCon, pcon, phoistAcyclic, plam, punsafeCoerce, (#), (:-->))
+import Plutarch (PlutusType (PInner, pcon', pmatch'), phoistAcyclic, plam, punsafeCoerce, (#), (:-->))
 import Plutarch.Internal (
   PType,
   RawTerm (RApply, RLamAbs, RVar),
@@ -28,9 +28,17 @@ type family ScottEncoded (r :: ((PType) -> Type) -> Type) (a :: PType) :: PType
 newtype ScottArgument r s t = ScottArgument {getScott :: Term s (ScottEncoded r t)}
 type ScottEncoding r t = ScottEncoded r t :--> t
 
-instance {-# OVERLAPS #-} Rank2.Foldable r => PCon (PRecord r) where
-  pcon :: forall s. PRecord r s -> Term s (PRecord r)
-  pcon = punsafeCoerce . rcon . getRecord
+instance (Rank2.Distributive r, Rank2.Traversable r) => PlutusType (PRecord r) where
+  type PInner (PRecord r) t = ScottEncoding r t
+  pcon' :: forall s. PRecord r s -> forall t. Term s (ScottEncoding r t)
+  pcon' (PRecord r) = rcon r
+  pmatch' :: forall s t. (forall t. Term s (ScottEncoding r t)) -> (PRecord r s -> Term s t) -> Term s t
+  pmatch' p f = p # arg
+    where
+      arg :: Term s (ScottEncoded r t)
+      arg = Term (\i -> TermResult (RLamAbs (fieldCount (initial @r) - 1) $ rawArg i) [])
+      rawArg :: Natural -> RawTerm
+      rawArg depth = getTerm $ asRawTerm (f $ PRecord variables) $ depth + fieldCount (initial @r)
 
 rcon :: forall r s t. Rank2.Foldable r => r (Term s) -> Term s (ScottEncoding r t)
 rcon r = plam (\f -> punsafeCoerce $ appEndo (getDual $ Rank2.foldMap (Dual . Endo . applyField) r) f)
@@ -50,9 +58,9 @@ letrec r = Term term
         (Dual rawTerms, deps) = Rank2.foldMap (rawResult . ($ n) . asRawTerm) (r selfReferring)
     rawResult TermResult {getTerm, getDeps} = (Dual [getTerm], getDeps)
     selfReferring = Rank2.fmap fromRecord accessors
-    fromRecord (ScottArgument (Term access)) = Term $ \depth -> mapTerm (\field -> RApply (RVar $ fieldCount + depth - 1) [field]) (access 0)
-    fieldCount :: Natural
-    fieldCount = getSum (Rank2.foldMap (const $ Sum 1) (accessors @r))
+    fromRecord :: ScottArgument r s a -> Term s a
+    fromRecord (ScottArgument (Term access)) =
+      Term $ \depth -> mapTerm (\field -> RApply (RVar $ fieldCount (initial @r) + depth - 1) [field]) (access 0)
 
 -- | Converts a Haskell field function to a Scott-encoded record field accessor.
 field ::
@@ -64,31 +72,39 @@ field f = getScott (f accessors)
 
 -- | Provides a record of function terms that access each field out of a Scott-encoded record.
 accessors :: forall r s. (Rank2.Distributive r, Rank2.Traversable r) => r (ScottArgument r s)
-accessors = Rank2.cotraverse accessor id
+accessors = abstract Rank2.<$> variables
   where
-    accessor :: (r (ScottArgument r s) -> ScottArgument r s a) -> ScottArgument r s a
-    accessor ref = ref ordered
-    ordered :: r (ScottArgument r s)
-    ordered = evalState (Rank2.traverse next initial) fieldCount
-    initial :: r (Compose Maybe (ScottArgument r s))
-    initial = Rank2.distribute Nothing
-    next :: f a -> State Natural (ScottArgument r s a)
+    abstract :: Term s a -> ScottArgument r s a
+    abstract (Term t) = ScottArgument (phoistAcyclic $ Term $ mapTerm (RLamAbs $ fieldCount (initial @r) - 1) . t)
+
+{- | A record of terms that each accesses a different variable in scope,
+ outside in following the field order.
+-}
+variables :: forall r s. (Rank2.Distributive r, Rank2.Traversable r) => r (Term s)
+variables = Rank2.cotraverse var id
+  where
+    var :: (r (Term s) -> Term s a) -> Term s a
+    var ref = ref ordered
+    ordered :: r (Term s)
+    ordered = evalState (Rank2.traverse next $ initial @r) (fieldCount $ initial @r)
+    next :: f a -> State Natural (Term s a)
     next _ = do
       i <- get
       let i' = pred i
       seq i' (put i')
-      return
-        ( ScottArgument $
-            phoistAcyclic $
-              Term $
-                const $
-                  TermResult
-                    { getTerm = RLamAbs (fieldCount - 1) $ RVar i'
-                    , getDeps = []
-                    }
-        )
-    fieldCount :: Natural
-    fieldCount = getSum (Rank2.foldMap (const $ Sum 1) initial)
+      return $
+        Term $
+          const $
+            TermResult
+              { getTerm = RVar i'
+              , getDeps = []
+              }
+
+initial :: Rank2.Distributive r => r (Compose Maybe (Term s))
+initial = Rank2.distribute Nothing
+
+fieldCount :: Rank2.Foldable r => r f -> Natural
+fieldCount = getSum . Rank2.foldMap (const $ Sum 1)
 
 -- | The raw Y-combinator term
 rfix :: RawTerm

diff --git a/Plutarch/Rec/TH.hs b/Plutarch/Rec/TH.hs
@@ -1,18 +1,23 @@
 {-# LANGUAGE TemplateHaskell #-}
 
-module Plutarch.Rec.TH (deriveScottEncoded) where
+module Plutarch.Rec.TH (deriveAll, deriveScottEncoded) where
 
 import Language.Haskell.TH (Q)
 import qualified Language.Haskell.TH as TH
 import Plutarch ((:-->))
 import Plutarch.Rec (ScottEncoded)
+import qualified Rank2.TH
 
+-- | Use as a TH splice for all necessary @instance@ declarations.
+deriveAll :: TH.Name -> Q [TH.Dec]
+deriveAll name = (<>) <$> deriveScottEncoded name <*> Rank2.TH.deriveAll name
+
+-- | Use as a TH splice for @type instance ScottEncoded@ declarations.
 deriveScottEncoded :: TH.Name -> Q [TH.Dec]
 deriveScottEncoded name = do
   con <- reifyConstructor name
   a <- TH.newName "a"
   let qa = pure (TH.VarT a)
-  --   _ <- [d| type instance ScottEncoded $(pure $ TH.ConT name) $qa = $(genScottEncoded con qa) |] >>= error . show
   [d|type instance ScottEncoded $(pure $ TH.ConT name) $qa = $(genScottEncoded con qa)|]
 
 genScottEncoded :: TH.Con -> Q TH.Type -> Q TH.Type

diff --git a/examples/Examples/LetRec.hs b/examples/Examples/LetRec.hs
@@ -2,12 +2,12 @@
 
 module Examples.LetRec (tests) where
 
-import Plutarch (printTerm, punsafeCoerce)
+import Plutarch (pcon', pmatch', printTerm)
 import Plutarch.Bool (PBool (PFalse, PTrue), pif, (#==))
 import Plutarch.Integer (PInteger)
 import Plutarch.Prelude
 import Plutarch.Rec (PRecord (PRecord), ScottEncoded, ScottEncoding, field, letrec)
-import Plutarch.Rec.TH (deriveScottEncoded)
+import Plutarch.Rec.TH (deriveAll)
 import Plutarch.String (PString)
 import qualified Rank2.TH
 import Test.Tasty (TestTree, testGroup)
@@ -21,25 +21,25 @@ data SampleRecord f = SampleRecord
   , sampleString :: f PString
   }
 
-sampleRecord :: PRecord SampleRecord s
-sampleRecord =
-  PRecord
-    SampleRecord
-      { sampleBool = pcon PFalse
-      , sampleInt = 6
-      , sampleString = "Salut, Monde!"
-      }
-
 data EvenOdd f = EvenOdd
   { even :: f (PInteger :--> PBool)
   , odd :: f (PInteger :--> PBool)
   }
 
 type instance ScottEncoded EvenOdd a = (PInteger :--> PBool) :--> (PInteger :--> PBool) :--> a
 
-$(deriveScottEncoded ''SampleRecord)
-$(Rank2.TH.deriveAll ''SampleRecord)
 $(Rank2.TH.deriveAll ''EvenOdd)
+$(deriveAll ''SampleRecord) -- also autoderives the @type instance ScottEncoded@
+
+sampleRecord :: Term (s :: S) (ScottEncoding SampleRecord (t :: PType))
+sampleRecord =
+  pcon' $
+    PRecord
+      SampleRecord
+        { sampleBool = pcon PFalse
+        , sampleInt = 6
+        , sampleString = "Salut, Monde!"
+        }
 
 sampleRecur :: Term (s :: S) (ScottEncoding SampleRecord (t :: PType))
 sampleRecur =
@@ -67,9 +67,13 @@ tests =
     "Records"
     [ testGroup
         "Simple"
-        [ testCase "precord" $
-            printTerm (punsafeCoerce (pcon sampleRecord) # field sampleInt)
+        [ testCase "record construction" $
+            printTerm (sampleRecord # field sampleInt)
               @?= "(program 1.0.0 ((\\i0 -> i1 False 6 \"Salut, Monde!\") (\\i0 -> \\i0 -> \\i0 -> i2)))"
+        , testCase "record field" $
+            equal' (sampleRecord # field sampleInt) "(program 1.0.0 6)"
+        , testCase "record match" $
+            equal' (pmatch' sampleRecord $ \(PRecord r) -> sampleString r) "(program 1.0.0 \"Salut, Monde!\")"
         ]
     , testGroup
         "Letrec"