@@ -1,124 +1,128 @@
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 Nock
14+import Data.Map.Strict (Map)
15+import Data.Map.Strict qualified as M
16+import GHC.Base (Word (..), Word#, int2Word#, isTrue#, leWord#, plusWord#, timesWord#)
17+import GHC.Int (Int (..))
18+import GHC.Num
1519import Nock.Types
16−import Numeric.Natural
17−
18−type BitParser = StateT (Int, Map Int Noun) BitGet
19−
20−parseBit :: BitParser Bool
21−parseBit = do
22− (counter, m) <- get
23− b <- lift getBool
24− put (counter + 1, m)
25− pure b
2620
27−parseBits :: Int -> BitParser [Bool]
28−parseBits 0 = pure []
29−parseBits n = do
30− b <- parseBit
31− (b :) <$> parseBits (n - 1)
32−
33−getOffset :: BitParser Int
21+type BitPrinterM = StateT (Natural, Map Noun Natural) Put.BitPut
22+
23+class (MonadState (Natural, Map Noun Natural) m) => BitPrinter m where
24+ putBit# :: (BitPrinter m) => Word# -> m ()
25+ putBool :: Bool -> m ()
26+
27+instance BitPrinter (StateT (Natural, Map Noun Natural) Identity) where
28+ putBit# _ = do
29+ (counter, m) <- get
30+ put (counter + 1, m)
31+ putBool _ = do
32+ (counter, m) <- get
33+ put (counter + 1, m)
34+
35+instance BitPrinter BitPrinterM where
36+ putBit# w = do
37+ (counter, m) <- get
38+ lift $ Put.putBool# w
39+ put (counter + 1, m)
40+ putBool b = do
41+ (counter, m) <- get
42+ lift $ Put.putBool b
43+ put (counter + 1, m)
44+
45+getOffset :: (BitPrinter m) => m Natural
3446getOffset = fst <$> get
3547
36−cue :: ByteString -> Noun
37−cue = runGet (runBitGet . withBitOrder LL . flip evalStateT (0, M.empty) $ cue')
48+getMapping :: (BitPrinter m) => m (Map Noun Natural)
49+getMapping = snd <$> get
3850
39−cue' :: BitParser Noun
40−cue' = do
41− offset <- getOffset
42− isNotAtom <- parseBit
43− case isNotAtom of
44− False -> do
45− a' <- rub'
46− let a = atom a'
47− (offset, m) <- get
48− put (offset, M.insert offset a m)
49− pure a
50− True -> do
51− isNotRef <- parseBit
52− case isNotRef of
53− False -> do
54− x <- cue'
55− y <- cue'
56− let c = cell x y
57− (offset, m) <- get
58− put (offset, M.insert offset c m)
59− pure c
60− True -> do
61− ref <- rub'
62− (_, m) <- get
63− case M.lookup (fromIntegral ref) m of
64− Nothing -> undefined
65− Just a -> pure a
51+putMapping :: (BitPrinter m) => Noun -> Natural -> m ()
52+putMapping idx n = do
53+ (offset, m) <- get
54+ put (offset, M.insert idx n m)
6655
6756jam :: Noun -> ByteString
68−jam n = runPut . runBitPut . withBitOrder LL $ jam' n
69−
70−jam' :: Noun -> BitPut ()
71−jam' (Atom n _) = do
72− putBool False
73− mat' n
74−jam' (Cell lhs rhs _) = do
75− putBool True
76− putBool False
77− jam' lhs
78− jam' rhs
57+jam n = runPut . Put.runBitPut . withBitOrder LL . flip evalStateT (0, M.empty) $ jam' n
58+
59+jam' :: (BitPrinter m) => Noun -> m ()
60+jam' n = do
61+ offset <- getOffset
62+ mapping <- getMapping
63+ case M.lookup n mapping of
64+ 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
75+ 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
7990
8091mat :: Natural -> ByteString
81−mat n = runPut . runBitPut . withBitOrder LL $ mat' n
92+mat n = runPut . Put.runBitPut . withBitOrder LL . flip evalStateT (0, M.empty) $ mat' n
8293
83−mat' :: Natural -> BitPut ()
84−mat' 0 = putBool True
94+mat' :: (BitPrinter m) => Natural -> m ()
95+mat' 0 = putBit# 1##
8596mat' n = do
86− let bits = numToBits n
87− bitsLength = length bits
88− lengthBits = init $ numToBits bitsLength
89− putBool False
97+ let bitsLength = countBits# n
98+ lengthBits = init $ numToBits (W# bitsLength)
99+ putBit# 0##
90100 putBits $ replicate (length lengthBits) False
91− putBool True
101+ putBit# 1##
92102 putBits lengthBits
93− putBits bits
103+ putNumToBits# n
104+
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
94111
95−putBits :: [Bool] -> BitPut ()
112+putBits :: (BitPrinter m) => [Bool] -> m ()
96113putBits [] = pure ()
97114putBits (x : xs) = putBool x >> putBits xs
98115
99−rub :: ByteString -> Natural
100−rub = runGet (runBitGet . withBitOrder LL . flip evalStateT (0, M.empty) $ rub')
101−
102−rub' :: BitParser Natural
103−rub' = do
104− parseBit >>= \case
105− True -> pure 0
106− False -> do
107− lengthOfLength <- countZeros 0
108− lengthBits <- (++ [True]) <$> parseBits lengthOfLength
109− let lent = bitsToNum lengthBits
110− bitsToNum <$> parseBits lent
111−
112−countZeros :: Int -> BitParser Int
113−countZeros n = do
114− b <- parseBit
115− case b of
116− True -> pure n
117− False -> countZeros (n + 1)
118−
119116numToBits :: (Bits a, Integral a) => a -> [Bool]
120117numToBits n | n <= 0 = []
121118numToBits n = ((n .&. 1) == 1) : numToBits (shiftR n 1)
122119
123−bitsToNum :: (Bits a, Integral a) => [Bool] -> a
124−bitsToNum = foldr (\b acc -> shiftL acc 1 .|. if b then 1 else 0) 0
120+countBits# :: Natural -> Word#
121+countBits# (NS 0##) = 0##
122+countBits# n = 1## `plusWord#` countBits# (n `naturalShiftR#` 1##)
123+
124+putNumToBits# :: (BitPrinter m) => Natural -> m ()
125+putNumToBits# (NS 0##) = pure ()
126+putNumToBits# n = case n .&. 1 of
127+ NS w -> putBit# w >> putNumToBits# (n `naturalShiftR#` 1##)
128+ _ -> undefined