@@ -11,20 +11,15 @@ 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
1514import Data.Map.Strict (Map)
1615import Data.Map.Strict qualified as M
17−import GHC.Base (Word (..), Word#, plusWord#)
16+import GHC.Base (Word (..), Word#, int2Word#, isTrue#, leWord#, plusWord#, timesWord#)
17+import GHC.Int (Int (..))
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−
2823class (MonadState (Natural, Map Noun Natural) m) => BitPrinter m where
2924 putBit# :: (BitPrinter m) => Word# -> m ()
3025 putBool :: Bool -> m ()
@@ -53,38 +48,45 @@ getOffset = fst <$> get
5348getMapping :: (BitPrinter m) => m (Map Noun Natural)
5449getMapping = snd <$> get
5550
56−putMapping :: (BitPrinter m) => Map Noun Natural -> m ()
57−putMapping m = do
58− (offset, _) <- get
59− put (offset, m)
51+putMapping :: (BitPrinter m) => Noun -> Natural -> m ()
52+putMapping idx n = do
53+ (offset, m) <- get
54+ put (offset, M.insert idx n m)
6055
6156jam :: Noun -> ByteString
6257jam n = runPut . Put.runBitPut . withBitOrder LL . flip evalStateT (0, M.empty) $ jam' n
6358
6459jam' :: (BitPrinter m) => Noun -> m ()
65−jam' (Atom n _) = do
66− putBit# 0##
67− mat' n
68−jam' n@(Cell lhs rhs _) = do
60+jam' n = do
6961 offset <- getOffset
70− putBit# 1##
71− let inlineAct = SomeBitPrinter $ do
72− putBit# 0##
73− jam' lhs
74− jam' rhs
7562 mapping <- getMapping
7663 case M.lookup n mapping of
7764 Nothing -> do
78− runSomeBitPrinter inlineAct
79− putMapping $ M.insert n offset mapping
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
8075 Just ref -> do
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
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
8890
8991mat :: Natural -> ByteString
9092mat n = runPut . Put.runBitPut . withBitOrder LL . flip evalStateT (0, M.empty) $ mat' n
@@ -100,6 +102,13 @@ mat' n = do
100102 putBits lengthBits
101103 putNumToBits# n
102104
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+
103112putBits :: (BitPrinter m) => [Bool] -> m ()
104113putBits [] = pure ()
105114putBits (x : xs) = putBool x >> putBits xs