@@ -1,136 +1,118 @@
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.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 (..))
1518import GHC.Num
16−import Nock
1719import Nock.Types
1820
19−type BitParser = StateT (Int, Map Int Noun) BitGet
20−
21−parseBit# :: BitParser Int
22−parseBit# = do
23− (counter, m) <- get
24− b <- lift getBool#
25− put (counter + 1, m)
26− pure b
27−
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)
34−
35−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
3646getOffset = fst <$> get
3747
38−getMapping :: BitParser (Map Int Noun)
48+getMapping :: (BitPrinter m) => m (Map Noun Natural)
3949getMapping = snd <$> get
4050
41−putMapping :: Map Int Noun -> BitParser ()
42−putMapping m = do
43− (offset, _) <- get
44− 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)
4555
46−cue :: ByteString -> Noun
47−cue = runGet (runBitGet . withBitOrder LL . flip evalStateT (0, M.empty) $ cue')
56+jam :: Noun -> ByteString
57+jam n = runPut . Put.runBitPut . withBitOrder LL . flip evalStateT (0, M.empty) $ jam' n
4858
49−cue' :: BitParser Noun
50−cue' = do
59+jam' :: (BitPrinter m) => Noun -> m ()
60+jam' n = do
5161 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−
80−jam :: Noun -> ByteString
81−jam n = runPut . runBitPut . withBitOrder LL $ jam' n
82−
83−jam' :: Noun -> BitPut ()
84−jam' (Atom n _) = do
85− putBool# 0##
86− mat' n
87−jam' (Cell lhs rhs _) = do
88− putBool# 1##
89− putBool# 0##
90− jam' lhs
91− jam' rhs
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
9290
9391mat :: Natural -> ByteString
94−mat n = runPut . runBitPut . withBitOrder LL $ mat' n
92+mat n = runPut . Put.runBitPut . withBitOrder LL . flip evalStateT (0, M.empty) $ mat' n
9593
96−mat' :: Natural -> BitPut ()
97−mat' 0 = putBool# 1##
94+mat' :: (BitPrinter m) => Natural -> m ()
95+mat' 0 = putBit# 1##
9896mat' n = do
9997 let bitsLength = countBits# n
10098 lengthBits = init $ numToBits (W# bitsLength)
101− putBool# 0##
99+ putBit# 0##
102100 putBits $ replicate (length lengthBits) False
103− putBool# 1##
101+ putBit# 1##
104102 putBits lengthBits
105103 putNumToBits# n
106104
107−putBits :: [Bool] -> BitPut ()
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+
112+putBits :: (BitPrinter m) => [Bool] -> m ()
108113putBits [] = pure ()
109114putBits (x : xs) = putBool x >> putBits xs
110115
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−
134116numToBits :: (Bits a, Integral a) => a -> [Bool]
135117numToBits n | n <= 0 = []
136118numToBits n = ((n .&. 1) == 1) : numToBits (shiftR n 1)
@@ -139,13 +121,8 @@ countBits# :: Natural -> Word#
139121countBits# (NS 0##) = 0##
140122countBits# n = 1## `plusWord#` countBits# (n `naturalShiftR#` 1##)
141123
142−putNumToBits# :: Natural -> BitPut ()
124+putNumToBits# :: (BitPrinter m) => Natural -> m ()
143125putNumToBits# (NS 0##) = pure ()
144126putNumToBits# n = case n .&. 1 of
145− NS w -> putBool# w >> putNumToBits# (n `naturalShiftR#` 1##)
127+ NS w -> putBit# w >> putNumToBits# (n `naturalShiftR#` 1##)
146128 _ -> 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