@@ -1,136 +1,109 @@
1−module Nock.Jam (jam, cue, mat, rub) where
1+{-# LANGUAGE NoMonomorphismRestriction #-}
22
3+module Nock.Jam (jam, mat) where
4+
5+import Control.Monad.Identity
6+import Control.Monad.State.Strict (MonadState (..))
37import Control.Monad.Trans.Class
4−import Control.Monad.Trans.State
8+import Control.Monad.Trans.State hiding (get, put)
59import Data.Binary.Bits.BitOrder
6−import Data.Binary.Bits.Get
7−import Data.Binary.Bits.Put
8−import Data.Binary.Get (runGet)
10+import Data.Binary.Bits.Put qualified as Put
911import Data.Binary.Put (runPut)
1012import Data.Bits
1113import Data.ByteString.Lazy (ByteString)
12−import Data.Map (Map)
13−import Data.Map qualified as M
14−import GHC.Base (Int (..), Word (..), Word#, int2Word#, minusWord#, or#, plusWord#, shiftL#, word2Int#)
14+import Data.HashMap.Strict qualified as HM
15+import Data.Map.Strict (Map)
16+import Data.Map.Strict qualified as M
17+import GHC.Base (Word (..), Word#, plusWord#)
1518import GHC.Num
16−import Nock
1719import Nock.Types
1820
19−type BitParser = StateT (Int, Map Int Noun) BitGet
21+type BitPrinterM = StateT (Natural, Map Noun Natural) Put.BitPut
22+
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
2027
21−parseBit# :: BitParser Int
22−parseBit# = do
23− (counter, m) <- get
24− b <- lift getBool#
25− put (counter + 1, m)
26− pure b
28+class (MonadState (Natural, Map Noun Natural) m) => BitPrinter m where
29+ putBit# :: (BitPrinter m) => Word# -> m ()
30+ putBool :: Bool -> m ()
2731
28−parseBits# :: Word# -> BitParser Natural
29−parseBits# 0## = pure 0
30−parseBits# n = do
31− I# b <- parseBit#
32− rest <- parseBits# (n `minusWord#` 1##)
33− pure $ (rest `shiftL` 1) `naturalOr` NS (int2Word# b)
32+instance BitPrinter (StateT (Natural, Map Noun Natural) Identity) where
33+ putBit# _ = do
34+ (counter, m) <- get
35+ put (counter + 1, m)
36+ putBool _ = do
37+ (counter, m) <- get
38+ put (counter + 1, m)
3439
35−getOffset :: BitParser Int
40+instance BitPrinter BitPrinterM where
41+ putBit# w = do
42+ (counter, m) <- get
43+ lift $ Put.putBool# w
44+ put (counter + 1, m)
45+ putBool b = do
46+ (counter, m) <- get
47+ lift $ Put.putBool b
48+ put (counter + 1, m)
49+
50+getOffset :: (BitPrinter m) => m Natural
3651getOffset = fst <$> get
3752
38−getMapping :: BitParser (Map Int Noun)
53+getMapping :: (BitPrinter m) => m (Map Noun Natural)
3954getMapping = snd <$> get
4055
41−putMapping :: Map Int Noun -> BitParser ()
56+putMapping :: (BitPrinter m) => Map Noun Natural -> m ()
4257putMapping m = do
4358 (offset, _) <- get
4459 put (offset, m)
4560
46−cue :: ByteString -> Noun
47−cue = runGet (runBitGet . withBitOrder LL . flip evalStateT (0, M.empty) $ cue')
48−
49−cue' :: BitParser Noun
50−cue' = do
51− offset <- getOffset
52− I# isNotAtom <- parseBit#
53− case isNotAtom of
54− 0# -> do
55− a' <- rub'
56− let a = atom a'
57− m <- getMapping
58− putMapping $ M.insert offset a m
59− pure a
60− _ -> do
61− I# isRef <- parseBit#
62− case isRef of
63− 0# -> do
64− x <- cue'
65− y <- cue'
66− let c = cell x y
67− m <- getMapping
68− putMapping $ M.insert offset c m
69− pure c
70− _ -> do
71− ref <- rub'
72− let ref# = case ref of
73− NS w -> I# (word2Int# w)
74− _ -> undefined
75− m <- getMapping
76− case M.lookup ref# m of
77− Nothing -> undefined
78− Just a -> pure a
79−
8061jam :: Noun -> ByteString
81−jam n = runPut . runBitPut . withBitOrder LL $ jam' n
62+jam n = runPut . Put.runBitPut . withBitOrder LL . flip evalStateT (0, M.empty) $ jam' n
8263
83−jam' :: Noun -> BitPut ()
64+jam' :: (BitPrinter m) => Noun -> m ()
8465jam' (Atom n _) = do
85− putBool# 0##
66+ putBit# 0##
8667 mat' n
87−jam' (Cell lhs rhs _) = do
88− putBool# 1##
89− putBool# 0##
90− jam' lhs
91− jam' rhs
68+jam' n@(Cell lhs rhs _) = do
69+ offset <- getOffset
70+ putBit# 1##
71+ let inlineAct = SomeBitPrinter $ do
72+ putBit# 0##
73+ jam' lhs
74+ jam' rhs
75+ mapping <- getMapping
76+ case M.lookup n mapping of
77+ Nothing -> do
78+ runSomeBitPrinter inlineAct
79+ putMapping $ M.insert n offset mapping
80+ 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
9288
9389mat :: Natural -> ByteString
94−mat n = runPut . runBitPut . withBitOrder LL $ mat' n
90+mat n = runPut . Put.runBitPut . withBitOrder LL . flip evalStateT (0, M.empty) $ mat' n
9591
96−mat' :: Natural -> BitPut ()
97−mat' 0 = putBool# 1##
92+mat' :: (BitPrinter m) => Natural -> m ()
93+mat' 0 = putBit# 1##
9894mat' n = do
9995 let bitsLength = countBits# n
10096 lengthBits = init $ numToBits (W# bitsLength)
101− putBool# 0##
97+ putBit# 0##
10298 putBits $ replicate (length lengthBits) False
103− putBool# 1##
99+ putBit# 1##
104100 putBits lengthBits
105101 putNumToBits# n
106102
107−putBits :: [Bool] -> BitPut ()
103+putBits :: (BitPrinter m) => [Bool] -> m ()
108104putBits [] = pure ()
109105putBits (x : xs) = putBool x >> putBits xs
110106
111−rub :: ByteString -> Natural
112−rub = runGet (runBitGet . withBitOrder LL . flip evalStateT (0, M.empty) $ rub')
113−
114−rub' :: BitParser Natural
115−rub' = do
116− parseBit# >>= \case
117− I# 0# -> do
118− W# lengthOfLength <- countZeros 0
119− lengthBits <- parseBits# lengthOfLength
120− let lengthBits# = case lengthBits of
121− NS w -> w
122− _ -> undefined
123− let lent = lengthBits# `or#` (1## `shiftL#` word2Int# lengthOfLength)
124− parseBits# lent
125− _ -> pure 0
126−
127−countZeros :: Word -> BitParser Word
128−countZeros n = do
129− I# b <- parseBit#
130− case b of
131− 0# -> countZeros (n + 1)
132− _ -> pure n
133−
134107numToBits :: (Bits a, Integral a) => a -> [Bool]
135108numToBits n | n <= 0 = []
136109numToBits n = ((n .&. 1) == 1) : numToBits (shiftR n 1)
@@ -139,13 +112,8 @@ countBits# :: Natural -> Word#
139112countBits# (NS 0##) = 0##
140113countBits# n = 1## `plusWord#` countBits# (n `naturalShiftR#` 1##)
141114
142−putNumToBits# :: Natural -> BitPut ()
115+putNumToBits# :: (BitPrinter m) => Natural -> m ()
143116putNumToBits# (NS 0##) = pure ()
144117putNumToBits# n = case n .&. 1 of
145− NS w -> putBool# w >> putNumToBits# (n `naturalShiftR#` 1##)
118+ NS w -> putBit# w >> putNumToBits# (n `naturalShiftR#` 1##)
146119 _ -> undefined
147−
148−-- ((n .&. 1) == 1) : numToBits (shiftR n 1)
149−
150−bitsToNum :: (Bits a, Integral a) => [Bool] -> a
151−bitsToNum = foldr (\b acc -> shiftL acc 1 .|. if b then 1 else 0) 0