MonadHashCons

covenant.cabal

@@ -81,12 +81,12 @@ common bench-lang
 library
   import: lang
   exposed-modules:
+    Control.Monad.HashCons
     Covenant.ASG
     Covenant.Constant
     Covenant.Prim
 
   other-modules:
-    Covenant.Internal.ASGBuilder
     Covenant.Internal.ASGNode
 
   build-depends:
@@ -101,6 +101,7 @@ library
     quickcheck-instances ==0.3.32,
     quickcheck-transformer ==0.3.1.2,
     text >=2.1.1 && <2.2,
+    transformers >=0.6.1.0 && <0.7.0.0,
     vector ==0.13.2.0,
 
   hs-source-dirs: src

src/Control/Monad/HashCons.hs

@@ -0,0 +1,147 @@
+{-# LANGUAGE FunctionalDependencies #-}
+
+-- | Module: Control.Monad.HashCons
+--
+-- Provides a transformer, and a capability type class in the style of @mtl@,
+-- for hash consing. See the Covenant wiki for how this works.
+module Control.Monad.HashCons
+  ( -- * Transformer
+    HashConsT,
+    runHashConsT,
+    hashCons,
+
+    -- * Capability type class
+    MonadHashCons (..),
+  )
+where
+
+import Control.Monad.State.Strict
+  ( StateT,
+    get,
+    modify,
+    runStateT,
+  )
+import Control.Monad.Trans (MonadTrans (lift))
+import Control.Monad.Trans.Except (ExceptT)
+import Control.Monad.Trans.Maybe (MaybeT)
+import Control.Monad.Trans.RWS.CPS (RWST)
+import Control.Monad.Trans.Reader (ReaderT)
+import Control.Monad.Trans.Writer.CPS (WriterT)
+import Data.Bimap (Bimap)
+import Data.Bimap qualified as Bimap
+import Data.Kind (Type)
+
+-- | A transformer implementing hash consing capabilities, with references of
+-- type @r@ and referents of type @e@. It is assumed that values of type @e@
+-- contain values of type @r@ in their capacity as references, though this is
+-- not a requirement of this transformer.
+--
+-- = Important note
+--
+-- This implementation is not suitable for any @m@ that throws exceptions. This
+-- includes @IO@, @ST@ and anything stacked atop them. For the reasons why, see
+-- [here](https://github.com/haskell-effectful/effectful/blob/master/transformers.md#statet).
+--
+-- @since 1.0.0
+newtype HashConsT (r :: Type) (e :: Type) (m :: Type -> Type) (a :: Type)
+  = HashConsT (StateT (Bimap r e) m a)
+  deriving
+    ( -- | @since 1.0.0
+      Functor,
+      -- | @since 1.0.0
+      Applicative,
+      -- | @since 1.0.0
+      Monad
+    )
+    via (StateT (Bimap r e) m)
+  deriving
+    ( -- | @since 1.0.0
+      MonadTrans
+    )
+    via StateT (Bimap r e)
+
+-- | Execute the computation described, returning both the result and the unique
+-- pairings of @r@ and @e@ produced as part of it.
+--
+-- @since 1.0.0
+runHashConsT ::
+  forall (r :: Type) (e :: Type) (m :: Type -> Type) (a :: Type).
+  HashConsT r e m a ->
+  m (a, Bimap r e)
+runHashConsT (HashConsT comp) = runStateT comp Bimap.empty
+
+-- | Given a value of type @e@, produce the unique value of type @r@ acting as a
+-- reference to it. This @r@ will be cached, ensuring any future requests for
+-- the reference for this value of type @e@ will be the same.
+--
+-- @since 1.0.0
+hashCons ::
+  forall (r :: Type) (e :: Type) (m :: Type -> Type).
+  (Ord r, Ord e, Bounded r, Enum r, Monad m) =>
+  e ->
+  HashConsT r e m r
+hashCons x = HashConsT $ do
+  binds <- get
+  case Bimap.lookupR x binds of
+    Nothing ->
+      if Bimap.null binds
+        then minBound <$ modify (Bimap.insert minBound x)
+        else do
+          let largestOldRef = fst . Bimap.findMax $ binds
+          let newRef = succ largestOldRef
+          newRef <$ modify (Bimap.insert newRef x)
+    Just ref -> pure ref
+
+-- | An @mtl@-style capability type class for hash consing capability, using
+-- references of type @r@ and values of type @e@.
+--
+-- = Laws
+--
+-- 1. @'refTo' x '>>' 'refTo' x@ @=@ @'refTo' x@
+--
+-- @since 1.0.0
+class
+  (Eq e, Eq r, Monad m) =>
+  MonadHashCons (r :: Type) (e :: Type) (m :: Type -> Type)
+    | m -> e r
+  where
+  -- | Produce the unique value of type @r@ that acts as a reference for the
+  -- given value of type @e@.
+  --
+  -- @since 1.0.0
+  refTo :: e -> m r
+
+-- | @since 1.0.0
+instance (Ord r, Ord e, Bounded r, Enum r, Monad m) => MonadHashCons r e (HashConsT r e m) where
+  {-# INLINEABLE refTo #-}
+  refTo = hashCons
+
+-- | @since 1.0.0
+instance (MonadHashCons r e m) => MonadHashCons r e (MaybeT m) where
+  {-# INLINEABLE refTo #-}
+  refTo e = lift (refTo e)
+
+-- | @since 1.0.0
+instance (MonadHashCons r e m) => MonadHashCons r e (ReaderT r' m) where
+  {-# INLINEABLE refTo #-}
+  refTo e = lift (refTo e)
+
+-- | @since 1.0.0
+instance (MonadHashCons r e m) => MonadHashCons r e (StateT s m) where
+  {-# INLINEABLE refTo #-}
+  refTo e = lift (refTo e)
+
+-- | @since 1.0.0
+instance (MonadHashCons r e m) => MonadHashCons r e (WriterT w m) where
+  {-# INLINEABLE refTo #-}
+  refTo e = lift (refTo e)
+
+-- | @since 1.0.0
+instance (MonadHashCons r e m) => MonadHashCons r e (RWST r' w s m) where
+  {-# INLINEABLE refTo #-}
+  refTo e = lift (refTo e)
+
+-- | @since 1.0.0
+instance (MonadHashCons r e m) => MonadHashCons r e (ExceptT e m) where
+  {-# INLINEABLE refTo #-}
+  refTo e = lift (refTo e)

src/Covenant/ASG.hs

@@ -24,7 +24,6 @@ module Covenant.ASG
     Ref (..),
     PrimCall (..),
     ASGNode,
-    ASGBuilder,
     Scope,
     ASG,
 
@@ -51,15 +50,8 @@ import Algebra.Graph.Acyclic.AdjacencyMap
     vertex,
   )
 import Algebra.Graph.AdjacencyMap qualified as Cyclic
-import Control.Monad.State.Strict (runState)
-import Covenant.Internal.ASGBuilder
-  ( ASGBuilder (ASGBuilder),
-    ASGBuilderState (ASGBuilderState),
-    app,
-    idOf,
-    lit,
-    prim,
-  )
+import Control.Monad.HashCons (HashConsT, MonadHashCons (refTo), runHashConsT)
+import Covenant.Constant (AConstant)
 import Covenant.Internal.ASGNode
   ( ASGNode (App, Lam, Let, Lit, Prim),
     Arg (Arg),
@@ -70,6 +62,7 @@ import Covenant.Internal.ASGNode
   )
 import Data.Bimap (Bimap)
 import Data.Bimap qualified as Bimap
+import Data.Kind (Type)
 import Data.Maybe (mapMaybe)
 import Data.Proxy (Proxy (Proxy))
 import GHC.TypeNats (CmpNat, KnownNat, natVal, type (+))
@@ -90,21 +83,26 @@ data ASG = ASG (Id, ASGNode) (AdjacencyMap (Id, ASGNode))
 -- refers to into a call graph. This is guaranteed to be acyclic.
 --
 -- @since 1.0.0
-toASG :: ASGBuilder Id -> Maybe ASG
-toASG (ASGBuilder comp) = do
-  let (start, ASGBuilderState binds) = runState comp (ASGBuilderState Bimap.empty)
-  if Bimap.size binds == 1
-    then do
-      -- This cannot fail, but the type system can't show it
-      initial <- (start,) <$> Bimap.lookup start binds
-      pure . ASG initial . vertex $ initial
-    else do
-      let asGraph = Cyclic.edges . go binds $ start
-      -- This cannot fail, but the type system can't show it
-      acyclic <- toAcyclic asGraph
-      -- Same as above
-      initial <- (start,) <$> Bimap.lookup start binds
-      pure . ASG initial $ acyclic
+toASG ::
+  forall (m :: Type -> Type).
+  (Monad m) =>
+  HashConsT Id ASGNode m Id ->
+  m (Maybe ASG)
+toASG comp = do
+  (start, binds) <- runHashConsT comp
+  pure $
+    if Bimap.size binds == 1
+      then do
+        -- This cannot fail, but the type system can't show it
+        initial <- (start,) <$> Bimap.lookup start binds
+        pure . ASG initial . vertex $ initial
+      else do
+        let asGraph = Cyclic.edges . go binds $ start
+        -- This cannot fail, but the type system can't show it
+        acyclic <- toAcyclic asGraph
+        -- Same as above
+        initial <- (start,) <$> Bimap.lookup start binds
+        pure . ASG initial $ acyclic
   where
     go ::
       Bimap Id ASGNode ->
@@ -139,6 +137,40 @@ data Scope (args :: Natural) (lets :: Natural) = Scope
 emptyScope :: Scope 0 0
 emptyScope = Scope
 
+-- | Construct a literal (constant) value.
+--
+-- @since 1.0.0
+lit ::
+  forall (m :: Type -> Type).
+  (MonadHashCons Id ASGNode m) =>
+  AConstant -> m Id
+lit = refTo . Lit
+
+-- | Construct a primitive function call.
+--
+-- @since 1.0.0
+prim ::
+  forall (m :: Type -> Type).
+  (MonadHashCons Id ASGNode m) =>
+  PrimCall -> m Id
+prim = refTo . Prim
+
+-- | Construct a function application. The first argument is (an expression
+-- evaluating to) a function, the second argument is (an expression evaluating
+-- to) an argument.
+--
+-- = Important note
+--
+-- Currently, this does not verify that the first argument is indeed a function,
+-- nor that the second argument is appropriate.
+--
+-- @since 1.0.0
+app ::
+  forall (m :: Type -> Type).
+  (MonadHashCons Id ASGNode m) =>
+  Ref -> Ref -> m Id
+app f x = refTo (App f x)
+
 -- | Given a proof of scope, construct one of the arguments in that scope. This
 -- requires use of @TypeApplications@ to select which argument you are
 -- interested in: argument @0@ is the one from the nearest lambda, argument @1@
@@ -184,13 +216,14 @@ bound Scope = Bound . fromIntegral . natVal $ Proxy @n
 --
 -- @since 1.0.0
 lam ::
-  forall (n :: Natural) (m :: Natural).
-  Scope n m ->
-  (Scope (n + 1) m -> ASGBuilder Ref) ->
-  ASGBuilder Id
+  forall (args :: Natural) (binds :: Natural) (m :: Type -> Type).
+  (MonadHashCons Id ASGNode m) =>
+  Scope args binds ->
+  (Scope (args + 1) binds -> m Ref) ->
+  m Id
 lam Scope f = do
   res <- f Scope
-  idOf . Lam $ res
+  refTo . Lam $ res
 
 -- | Given a proof of scope, a 'Ref' to an expression to bind to, and a function
 -- to construct a @let@-binding body using a \'larger\' proof of scope, construct
@@ -208,14 +241,15 @@ lam Scope f = do
 --
 -- @since 1.0.0
 letBind ::
-  forall (n :: Natural) (m :: Natural).
-  Scope n m ->
+  forall (args :: Natural) (binds :: Natural) (m :: Type -> Type).
+  (MonadHashCons Id ASGNode m) =>
+  Scope args binds ->
   Ref ->
-  (Scope n (m + 1) -> ASGBuilder Ref) ->
-  ASGBuilder Id
+  (Scope args (binds + 1) -> m Ref) ->
+  m Id
 letBind Scope r f = do
   res <- f Scope
-  idOf . Let r $ res
+  refTo . Let r $ res
 
 -- Helpers