@@ -11,56 +11,69 @@ import Data.Bits
1111import Data.ByteString.Lazy (ByteString)
1212import Data.Map (Map)
1313import Data.Map qualified as M
14+import GHC.Base (Int (..), Word (..), Word#, int2Word#, minusWord#, or#, plusWord#, shiftL#, word2Int#)
15+import GHC.Num
1416import Nock
1517import Nock.Types
16−import Numeric.Natural
1718
1819type BitParser = StateT (Int, Map Int Noun) BitGet
1920
20−parseBit :: BitParser Bool
21−parseBit = do
21+parseBit# :: BitParser Int
22+parseBit# = do
2223 (counter, m) <- get
23− b <- lift getBool
24+ b <- lift getBool#
2425 put (counter + 1, m)
2526 pure b
2627
27−parseBits :: Int -> BitParser [Bool]
28−parseBits 0 = pure []
29−parseBits n = do
30− b <- parseBit
31− (b :) <$> parseBits (n - 1)
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)
3234
3335getOffset :: BitParser Int
3436getOffset = fst <$> get
3537
38+getMapping :: BitParser (Map Int Noun)
39+getMapping = snd <$> get
40+
41+putMapping :: Map Int Noun -> BitParser ()
42+putMapping m = do
43+ (offset, _) <- get
44+ put (offset, m)
45+
3646cue :: ByteString -> Noun
3747cue = runGet (runBitGet . withBitOrder LL . flip evalStateT (0, M.empty) $ cue')
3848
3949cue' :: BitParser Noun
4050cue' = do
4151 offset <- getOffset
42− isNotAtom <- parseBit
52+ I# isNotAtom <- parseBit#
4353 case isNotAtom of
44− False -> do
54+ 0# -> do
4555 a' <- rub'
4656 let a = atom a'
47− (offset, m) <- get
48− put (offset, M.insert offset a m)
57+ m <- getMapping
58+ putMapping $ M.insert offset a m
4959 pure a
50− True -> do
51− isNotRef <- parseBit
52− case isNotRef of
53− False -> do
60+ _ -> do
61+ I# isRef <- parseBit#
62+ case isRef of
63+ 0# -> do
5464 x <- cue'
5565 y <- cue'
5666 let c = cell x y
57− (offset, m) <- get
58− put (offset, M.insert offset c m)
67+ m <- getMapping
68+ putMapping $ M.insert offset c m
5969 pure c
60− True -> do
70+ _ -> do
6171 ref <- rub'
62− (_, m) <- get
63− case M.lookup (fromIntegral ref) m of
72+ let ref# = case ref of
73+ NS w -> I# (word2Int# w)
74+ _ -> undefined
75+ m <- getMapping
76+ case M.lookup ref# m of
6477 Nothing -> undefined
6578 Just a -> pure a
6679
@@ -69,11 +82,11 @@ jam n = runPut . runBitPut . withBitOrder LL $ jam' n
6982
7083jam' :: Noun -> BitPut ()
7184jam' (Atom n _) = do
72− putBool False
85+ putBool# 0##
7386 mat' n
7487jam' (Cell lhs rhs _) = do
75− putBool True
76− putBool False
88+ putBool# 1##
89+ putBool# 0##
7790 jam' lhs
7891 jam' rhs
7992
@@ -81,16 +94,15 @@ mat :: Natural -> ByteString
8194mat n = runPut . runBitPut . withBitOrder LL $ mat' n
8295
8396mat' :: Natural -> BitPut ()
84−mat' 0 = putBool True
97+mat' 0 = putBool# 1##
8598mat' n = do
86− let bits = numToBits n
87− bitsLength = length bits
88− lengthBits = init $ numToBits bitsLength
89− putBool False
99+ let bitsLength = countBits# n
100+ lengthBits = init $ numToBits (W# bitsLength)
101+ putBool# 0##
90102 putBits $ replicate (length lengthBits) False
91− putBool True
103+ putBool# 1##
92104 putBits lengthBits
93− putBits bits
105+ putNumToBits# n
94106
95107putBits :: [Bool] -> BitPut ()
96108putBits [] = pure ()
@@ -101,24 +113,39 @@ rub = runGet (runBitGet . withBitOrder LL . flip evalStateT (0, M.empty) $ rub')
101113
102114rub' :: BitParser Natural
103115rub' = 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
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
113128countZeros n = do
114− b <- parseBit
129+ I# b <- parseBit#
115130 case b of
116− True -> pure n
117− False -> countZeros (n + 1)
131+ 0# -> countZeros (n + 1)
132+ _ -> pure n
118133
119134numToBits :: (Bits a, Integral a) => a -> [Bool]
120135numToBits n | n <= 0 = []
121136numToBits n = ((n .&. 1) == 1) : numToBits (shiftR n 1)
122137
138+countBits# :: Natural -> Word#
139+countBits# (NS 0##) = 0##
140+countBits# n = 1## `plusWord#` countBits# (n `naturalShiftR#` 1##)
141+
142+putNumToBits# :: Natural -> BitPut ()
143+putNumToBits# (NS 0##) = pure ()
144+putNumToBits# n = case n .&. 1 of
145+ NS w -> putBool# w >> putNumToBits# (n `naturalShiftR#` 1##)
146+ _ -> undefined
147+
148+-- ((n .&. 1) == 1) : numToBits (shiftR n 1)
149+
123150bitsToNum :: (Bits a, Integral a) => [Bool] -> a
124151bitsToNum = foldr (\b acc -> shiftL acc 1 .|. if b then 1 else 0) 0