@@ -1,118 +1,136 @@
1−{-# LANGUAGE NoMonomorphismRestriction #-}
1+module Nock.Jam (jam, cue, mat, rub) where
22
3−module Nock.Jam (jam, mat) where
4−
5−import Control.Monad.Identity
6−import Control.Monad.State.Strict (MonadState (..))
73import Control.Monad.Trans.Class
8−import Control.Monad.Trans.State hiding (get, put)
4+import Control.Monad.Trans.State
95import Data.Binary.Bits.BitOrder
10−import Data.Binary.Bits.Put qualified as Put
6+import Data.Binary.Bits.Get
7+import Data.Binary.Bits.Put
8+import Data.Binary.Get (runGet)
119import Data.Binary.Put (runPut)
1210import Data.Bits
1311import Data.ByteString.Lazy (ByteString)
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 (..))
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#)
1815import GHC.Num
16+import Nock
1917import Nock.Types
2018
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
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
4636getOffset = fst <$> get
4737
48−getMapping :: (BitPrinter m) => m (Map Noun Natural)
38+getMapping :: BitParser (Map Int Noun)
4939getMapping = snd <$> get
5040
51−putMapping :: (BitPrinter m) => Noun -> Natural -> m ()
52−putMapping idx n = do
53− (offset, m) <- get
54− put (offset, M.insert idx n m)
41+putMapping :: Map Int Noun -> BitParser ()
42+putMapping m = do
43+ (offset, _) <- get
44+ put (offset, m)
5545
56−jam :: Noun -> ByteString
57−jam n = runPut . Put.runBitPut . withBitOrder LL . flip evalStateT (0, M.empty) $ jam' n
46+cue :: ByteString -> Noun
47+cue = runGet (runBitGet . withBitOrder LL . flip evalStateT (0, M.empty) $ cue')
5848
59−jam' :: (BitPrinter m) => Noun -> m ()
60−jam' n = do
49+cue' :: BitParser Noun
50+cue' = do
6151 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
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
9092
9193mat :: Natural -> ByteString
92−mat n = runPut . Put.runBitPut . withBitOrder LL . flip evalStateT (0, M.empty) $ mat' n
94+mat n = runPut . runBitPut . withBitOrder LL $ mat' n
9395
94−mat' :: (BitPrinter m) => Natural -> m ()
95−mat' 0 = putBit# 1##
96+mat' :: Natural -> BitPut ()
97+mat' 0 = putBool# 1##
9698mat' n = do
9799 let bitsLength = countBits# n
98100 lengthBits = init $ numToBits (W# bitsLength)
99− putBit# 0##
101+ putBool# 0##
100102 putBits $ replicate (length lengthBits) False
101− putBit# 1##
103+ putBool# 1##
102104 putBits lengthBits
103105 putNumToBits# n
104106
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 ()
107+putBits :: [Bool] -> BitPut ()
113108putBits [] = pure ()
114109putBits (x : xs) = putBool x >> putBits xs
115110
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+
116134numToBits :: (Bits a, Integral a) => a -> [Bool]
117135numToBits n | n <= 0 = []
118136numToBits n = ((n .&. 1) == 1) : numToBits (shiftR n 1)
@@ -121,8 +139,13 @@ countBits# :: Natural -> Word#
121139countBits# (NS 0##) = 0##
122140countBits# n = 1## `plusWord#` countBits# (n `naturalShiftR#` 1##)
123141
124−putNumToBits# :: (BitPrinter m) => Natural -> m ()
142+putNumToBits# :: Natural -> BitPut ()
125143putNumToBits# (NS 0##) = pure ()
126144putNumToBits# n = case n .&. 1 of
127− NS w -> putBit# w >> putNumToBits# (n `naturalShiftR#` 1##)
145+ NS w -> putBool# w >> putNumToBits# (n `naturalShiftR#` 1##)
128146 _ -> 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