Serialising data structures with sharing (i.e. after "tying the knot")

SerialiseKnots.hs

{-# LANGUAGE ExistentialQuantification #-}

module SerialiseKnots where

import System.Mem.StableName
import Data.Map (Map)
import qualified Data.Map as Map
import Control.Monad.Trans.State
import Control.Monad.Trans.Class
import Control.Monad (when)

data Ref = forall a. Ref (StableName a)

instance Eq Ref where
  (==) (Ref a) (Ref b) = eqStableName a b

instance Ord Ref where
  -- HACK. Hashing can give false positives
  compare (Ref a) (Ref b) = compare (hashStableName a) (hashStableName b)

type Ser a = StateT (Map Ref String) IO a

class SerialiseKnots a where
  serialiseKnots :: a -> Ser String

instance SerialiseKnots Int where
  serialiseKnots a = pure (show a)

data List a = Empty | Node a (List a)

instance SerialiseKnots a => SerialiseKnots (List a) where
  serialiseKnots Empty = pure "Empty"
  serialiseKnots (Node x xs) = do
    ref <- fmap Ref (lift $ makeStableName x)
    refs <- get
    when (not $ Map.member ref refs) $ do
      rep <- serialiseKnots x
      refs <- get
      put $ Map.insert ref rep refs
    rest <- serialiseKnots xs
    pure $ makeVar ref ++ ":" ++ rest

makeVar :: Ref -> String
makeVar (Ref s) = "v" ++ show (hashStableName s)


serialise :: SerialiseKnots a => a -> IO String
serialise a = do
  (rep, refs) <- runStateT (serialiseKnots a) Map.empty
  pure $ "let\n" ++ unlines (serialiseIndex refs) ++ "\nin\n" ++ rep

serialiseIndex :: Map Ref String -> [String]
serialiseIndex = Map.foldrWithKey f []
  where
    f ref v s = (makeVar ref ++ " = " ++ v):s


sample :: List (List Int)
sample = Node x (Node x (Node x Empty))
  where
    x = Node 12 Empty

main :: IO ()
main = do
  v <- serialise sample
  putStrLn v

-- Sample run:
--
-- λ stack run
-- let
-- v1 = v2:Empty
-- v2 = 12
--
-- in
-- v1:v1:v1:Empty