@@ -11,15 +11,20 @@ import Data.Binary.Bits.Put qualified as Put
1111import Data.Binary.Put (runPut)
1212import Data.Bits
1313import Data.ByteString.Lazy (ByteString)
14+import Data.HashMap.Strict qualified as HM
1415import Data.Map.Strict (Map)
1516import Data.Map.Strict qualified as M
16−import GHC.Base (Word (..), Word#, int2Word#, isTrue#, leWord#, plusWord#, timesWord#)
17−import GHC.Int (Int (..))
17+import GHC.Base (Word (..), Word#, plusWord#)
1818import GHC.Num
1919import Nock.Types
2020
2121type BitPrinterM = StateT (Natural, Map Noun Natural) Put.BitPut
2222
23+newtype SomeBitPrinter a = SomeBitPrinter (forall m. (BitPrinter m) => m a)
24+
25+runSomeBitPrinter :: (BitPrinter m) => SomeBitPrinter a -> m a
26+runSomeBitPrinter (SomeBitPrinter act) = act
27+
2328class (MonadState (Natural, Map Noun Natural) m) => BitPrinter m where
2429 putBit# :: (BitPrinter m) => Word# -> m ()
2530 putBool :: Bool -> m ()
@@ -48,45 +53,38 @@ getOffset = fst <$> get
4853getMapping :: (BitPrinter m) => m (Map Noun Natural)
4954getMapping = snd <$> get
5055
51−putMapping :: (BitPrinter m) => Noun -> Natural -> m ()
52−putMapping idx n = do
53− (offset, m) <- get
54− put (offset, M.insert idx n m)
56+putMapping :: (BitPrinter m) => Map Noun Natural -> m ()
57+putMapping m = do
58+ (offset, _) <- get
59+ put (offset, m)
5560
5661jam :: Noun -> ByteString
5762jam n = runPut . Put.runBitPut . withBitOrder LL . flip evalStateT (0, M.empty) $ jam' n
5863
5964jam' :: (BitPrinter m) => Noun -> m ()
60−jam' n = do
65+jam' (Atom n _) = do
66+ putBit# 0##
67+ mat' n
68+jam' n@(Cell lhs rhs _) = do
6169 offset <- getOffset
70+ putBit# 1##
71+ let inlineAct = SomeBitPrinter $ do
72+ putBit# 0##
73+ jam' lhs
74+ jam' rhs
6275 mapping <- getMapping
6376 case M.lookup n mapping of
6477 Nothing -> do
65− putMapping n offset
66− case n of
67− (Atom a _) -> do
68− putBit# 0##
69− mat' a
70− (Cell lhs rhs _) -> do
71− putBit# 1##
72− putBit# 0##
73− jam' lhs
74− jam' rhs
78+ runSomeBitPrinter inlineAct
79+ putMapping $ M.insert n offset mapping
7580 Just ref -> do
76− case n of
77− Cell {} -> do
78− putBit# 1##
79− putBit# 1##
80− mat' ref
81− (Atom a _) -> do
82− if isTrue# (matSize a `leWord#` matSize ref)
83− then do
84− putBit# 0##
85− mat' a
86− else do
87− putBit# 1##
88− putBit# 1##
89− mat' ref
81+ s <- get
82+ let refAct = SomeBitPrinter $ do
83+ putBit# 1##
84+ mat' ref
85+ (inlineLength, _) = runIdentity . flip execStateT s . runSomeBitPrinter $ inlineAct
86+ (refLength, _) = runIdentity . flip execStateT s . runSomeBitPrinter $ refAct
87+ if inlineLength < refLength then runSomeBitPrinter inlineAct else runSomeBitPrinter refAct
9088
9189mat :: Natural -> ByteString
9290mat n = runPut . Put.runBitPut . withBitOrder LL . flip evalStateT (0, M.empty) $ mat' n
@@ -102,13 +100,6 @@ mat' n = do
102100 putBits lengthBits
103101 putNumToBits# n
104102
105−matSize :: Natural -> Word#
106−matSize 0 = 1##
107−matSize n =
108− let bitsLength = countBits# n
109− !(I# lengthBits) = length $ init $ numToBits (W# bitsLength)
110− in 2## `plusWord#` (int2Word# lengthBits `timesWord#` 2##) `plusWord#` bitsLength
111−
112103putBits :: (BitPrinter m) => [Bool] -> m ()
113104putBits [] = pure ()
114105putBits (x : xs) = putBool x >> putBits xs