mirror of
https://github.com/urbit/shrub.git
synced 2024-12-23 02:41:35 +03:00
219 lines
7.4 KiB
Haskell
219 lines
7.4 KiB
Haskell
module Data.Noun.Jam where
|
|
|
|
import ClassyPrelude
|
|
import Data.Noun
|
|
import Data.Noun.Atom
|
|
import Data.Noun.Poet
|
|
import Data.Bits
|
|
import Control.Lens
|
|
import Text.Printf
|
|
|
|
import Data.Map (Map)
|
|
import Control.Monad (guard)
|
|
|
|
import Test.Tasty
|
|
import Test.Tasty.TH
|
|
import Test.Tasty.QuickCheck as QC
|
|
import Test.QuickCheck
|
|
|
|
|
|
-- Length-Encoded Atoms --------------------------------------------------------
|
|
|
|
bex :: (Num a, Bits a) => Int -> a
|
|
bex = shiftL 1
|
|
|
|
mat' :: Atom -> Buf
|
|
mat' 0 = Buf 1 1
|
|
mat' atm = Buf bufWid buffer
|
|
where
|
|
atmWid = bitWidth atm
|
|
preWid = bitWidth (toAtom atmWid)
|
|
bufWid = preWid + preWid + atmWid - 1
|
|
prefix = bex preWid
|
|
extras = takeBits (preWid-1) (toAtom atmWid)
|
|
suffix = xor extras (shiftL (takeBits (atmWid-1) atm) (preWid-1))
|
|
buffer = bitConcat prefix suffix
|
|
|
|
rub' :: Cursor -> Maybe Buf
|
|
rub' slc@(Cursor idx buf) =
|
|
leadingZeros slc >>= \case
|
|
0 -> pure (Buf 1 0)
|
|
prefix -> pure (Buf sz val)
|
|
where
|
|
widIdx = idx + 1 + prefix
|
|
width = fromSlice (Slice widIdx (prefix - 1) buf)
|
|
datIdx = widIdx + (prefix-1)
|
|
datWid = fromIntegral (2^(prefix-1) + width) - 1
|
|
sz = datWid + (2*prefix)
|
|
val = bex datWid .|. fromSlice (Slice datIdx datWid buf)
|
|
|
|
jam' :: Noun -> Atom
|
|
jam' = toAtom . fst . go 0 mempty
|
|
where
|
|
insertNoun :: Noun -> Int -> Map Noun Int -> Map Noun Int
|
|
insertNoun n i tbl = lookup n tbl
|
|
& maybe tbl (const $ insertMap n i tbl)
|
|
|
|
go :: Int -> Map Noun Int -> Noun -> (Buf, Map Noun Int)
|
|
go off oldTbl noun =
|
|
let tbl = insertNoun noun off oldTbl in
|
|
case (lookup noun oldTbl, noun) of
|
|
(Just ref, Atom atm) | bitWidth atm <= bitWidth (toAtom ref) ->
|
|
(Buf (1+sz) (shiftL res 1), tbl)
|
|
where Buf sz res = mat' atm
|
|
(Just ref, _) ->
|
|
(Buf (2+sz) (xor 3 (shiftL res 2)), tbl)
|
|
where Buf sz res = mat' (toAtom ref)
|
|
(Nothing, Atom atm) ->
|
|
(Buf (1+sz) (shiftL res 1), tbl)
|
|
where Buf sz res = mat' atm
|
|
(Nothing, Cell lef rit) ->
|
|
(Buf (2+lSz+rSz) (xor 1 (shiftL (lRes .|. shiftL rRes lSz) 2)), rTbl)
|
|
where (Buf lSz lRes, lTbl) = go (off+2) tbl lef
|
|
(Buf rSz rRes, rTbl) = go (off+lSz) lTbl rit
|
|
|
|
cue' :: Atom -> Maybe Noun
|
|
cue' buf = view _2 <$> go mempty 0
|
|
where
|
|
go :: Map Int Noun -> Int -> Maybe (Int, Noun, Map Int Noun)
|
|
go tbl i =
|
|
case (bitIdx i buf, bitIdx (i+1) buf) of
|
|
(False, _ ) -> do Buf wid at <- rub' (Cursor (i+1) buf)
|
|
let r = toNoun at
|
|
pure (wid+1, r, insertMap i r tbl)
|
|
(True, False ) -> do (lSz,lef,tbl) <- go tbl (i+2)
|
|
(rSz,rit,tbl) <- go tbl (i+2+fromIntegral lSz)
|
|
let r = Cell lef rit
|
|
pure (2+lSz+rSz, r, insertMap i r tbl)
|
|
(True, True ) -> do Buf wid at <- rub' (Cursor (i+2) buf)
|
|
r <- lookup (fromIntegral at) tbl & \case
|
|
Nothing -> error ("bad-ref-" <> show at)
|
|
Just ix -> Just ix
|
|
pure (2+wid, r, tbl)
|
|
|
|
--------------------------------------------------------------------------------
|
|
|
|
mat :: Atom -> Buf
|
|
mat 0 = Buf 1 1
|
|
mat atm = Buf bufWid buffer
|
|
where
|
|
atmWid = bitWidth atm
|
|
preWid = bitWidth (toAtom atmWid)
|
|
bufWid = preWid + preWid + atmWid
|
|
prefix = shiftL 1 preWid
|
|
extras = takeBits (preWid-1) (toAtom atmWid)
|
|
suffix = xor extras (shiftL atm (preWid-1))
|
|
buffer = bitConcat prefix suffix
|
|
|
|
bufVal Nothing = "<nil>"
|
|
bufVal (Just (Buf sz v)) = show v <> " [" <> show sz <> "]"
|
|
|
|
rub :: Cursor -> Maybe Buf
|
|
rub slc@(Cursor idx buf) =
|
|
leadingZeros slc >>= \case
|
|
0 -> pure (Buf 1 0)
|
|
prefix -> pure (Buf sz val)
|
|
where
|
|
widIdx = idx + 1 + prefix
|
|
width = fromSlice (Slice widIdx (prefix - 1) buf)
|
|
datIdx = widIdx + (prefix-1)
|
|
datWid = fromIntegral $ 2^(prefix-1) + width
|
|
sz = datWid + (2*prefix)
|
|
val = fromSlice (Slice datIdx datWid buf)
|
|
|
|
-- Noun Serialization ----------------------------------------------------------
|
|
|
|
-- bex can be implemented using
|
|
-- `mpz_mul_2exp(a_mp, a_mp, a);
|
|
|
|
jam :: Noun -> Atom
|
|
jam = toAtom . fst . go 0 mempty
|
|
where
|
|
insertNoun :: Noun -> Int -> Map Noun Int -> Map Noun Int
|
|
insertNoun n i tbl = lookup n tbl
|
|
& maybe tbl (const $ insertMap n i tbl)
|
|
|
|
go :: Int -> Map Noun Int -> Noun -> (Buf, Map Noun Int)
|
|
go off oldTbl noun =
|
|
let tbl = insertNoun noun off oldTbl in
|
|
case (lookup noun oldTbl, noun) of
|
|
(Just ref, Atom atm) | bitWidth atm <= bitWidth (toAtom ref) ->
|
|
(Buf (1+sz) (shiftL res 1), tbl)
|
|
where Buf sz res = mat atm
|
|
(Just ref, _) ->
|
|
(Buf (2+sz) (xor 3 (shiftL res 2)), tbl)
|
|
where Buf sz res = mat (toAtom ref)
|
|
(Nothing, Atom atm) ->
|
|
(Buf (1+sz) (shiftL res 1), tbl)
|
|
where Buf sz res = mat atm
|
|
(Nothing, Cell lef rit) ->
|
|
(Buf (2+lSz+rSz) (xor 1 (shiftL (bitConcat lRes rRes) 2)), rTbl)
|
|
where (Buf lSz lRes, lTbl) = go (off+2) tbl lef
|
|
(Buf rSz rRes, rTbl) = go (off+lSz) lTbl rit
|
|
|
|
|
|
|
|
leadingZeros :: Cursor -> Maybe Int
|
|
leadingZeros (Cursor idx buf) = go 0
|
|
where wid = bitWidth buf
|
|
go n = do () <- if (n < wid) then pure ()
|
|
else error "infinite-atom"
|
|
guard (n < wid)
|
|
if bitIdx (idx+n) buf then pure n else go (n+1)
|
|
|
|
cue :: Atom -> Maybe Noun
|
|
cue buf = view _2 <$> go mempty 0
|
|
where
|
|
go :: Map Int Noun -> Int -> Maybe (Int, Noun, Map Int Noun)
|
|
go tbl i =
|
|
-- trace ("go-" <> show i)
|
|
case (bitIdx i buf, bitIdx (i+1) buf) of
|
|
(False, _ ) -> do Buf wid at <- rub (Cursor (i+1) buf)
|
|
let r = toNoun at
|
|
pure (wid+1, r, insertMap i r tbl)
|
|
(True, False ) -> do (lSz,lef,tbl) <- go tbl (i+2)
|
|
(rSz,rit,tbl) <- go tbl (i+2+fromIntegral lSz)
|
|
let r = Cell lef rit
|
|
pure (2+lSz+rSz, r, insertMap i r tbl)
|
|
(True, True ) -> do Buf wid at <- rub (Cursor (i+2) buf)
|
|
-- traceM ("ref-" <> show at)
|
|
r <- lookup (fromIntegral at) tbl & \case
|
|
Nothing -> error ("bad-ref-" <> show at)
|
|
Just ix -> Just ix
|
|
pure (2+wid, r, tbl)
|
|
|
|
|
|
-- Tests -----------------------------------------------------------------------
|
|
|
|
pills :: [Atom]
|
|
pills = [ 0x2, 0xc, 0x48, 0x29, 0xc9, 0x299
|
|
, 0x3170_c7c1, 0x93_c7c1, 0xa_72e0, 0x1bd5_b7dd_e080
|
|
]
|
|
|
|
cueTest :: Maybe [Noun]
|
|
cueTest = traverse cue pills
|
|
|
|
jamTest :: Maybe [Atom]
|
|
jamTest = fmap jam <$> cueTest
|
|
|
|
prop_jamCue :: Noun -> Bool
|
|
prop_jamCue n = Just n == cue (jam n)
|
|
|
|
prop_matRub :: Atom -> Bool
|
|
prop_matRub atm = matSz==rubSz && rubRes==atm
|
|
where
|
|
Buf matSz matBuf = mat atm
|
|
Buf rubSz rubRes = fromMaybe mempty (rub $ Cursor 0 matBuf)
|
|
|
|
prop_jamCue' :: Noun -> Bool
|
|
prop_jamCue' n = Just n == cue' (jam' n)
|
|
|
|
prop_matRub' :: Atom -> Bool
|
|
prop_matRub' atm = matSz==rubSz && rubRes==atm
|
|
where
|
|
Buf matSz matBuf = mat' atm
|
|
Buf rubSz rubRes = fromMaybe mempty (rub' $ Cursor 0 matBuf)
|
|
|
|
main :: IO ()
|
|
main = $(defaultMainGenerator)
|