packed-lists.hs

packed-lists.hs

#!/usr/bin/env nix-shell
#! nix-shell -p "ghc.withPackages (pkgs: with pkgs; [ bytestring vector lens ]) " ghcid
-- #! nix-shell -i "ghcid -c 'ghci -Wall' -T main"
-- #! nix-shell -i "ghcid"
#! nix-shell -i "runghc --ghc-arg=-threaded --ghc-arg=-rtsopts"

{-# LANGUAGE TemplateHaskell #-}

import Control.Monad.ST
import Debug.Trace
import Data.Foldable
import Data.Traversable
import Data.Word
import Data.Bits
import Control.Lens.Operators
import Data.Bits.Lens
import Data.Functor
import Control.Monad
import qualified Numeric.Lens as Num
import qualified Data.Vector as V
import Data.Vector (Vector)
import qualified Data.Vector.Storable.Mutable as SV
import Data.Vector.Storable.Mutable (MVector, PrimState, PrimMonad)

type BitPack s = MVector s Word8
data PackedList s
  = Nil
  | PackedCons {- UNPACK -} !(BitPack s) !(PackedList s)

-- We are allocating a size of 64 bytes.
-- If every int is 24 bits(3 bytes long), we can pack 64 `quot` 3 = 21, having one byte to spare.
-- The last byte tell where is the next position to insert.
-- if it is i, we insert the next number from i to i+2 and set the value to i+3(for the next insert)

idxByte :: Int
idxByte = 63

maxNodeCapacity :: Word8 -- If this is the contents of idxByte, then this node is full
maxNodeCapacity = 63

new :: PrimMonad m => m (BitPack (PrimState m))
new = SV.replicate 64 zeroBits

cons :: PrimMonad m => Word32 -> PackedList (PrimState m) -> m (PackedList (PrimState m))
cons i packedList = do
  (insertableHd, tl) <-
    case packedList of
      Nil -> (,Nil) <$> new
      pl@(PackedCons hd tl) -> do
        isFull <- nodeFull hd
        if isFull
          then (,pl) <$> new
          else return $! (hd, tl)
  insertInPackedBS i insertableHd
  return $! PackedCons insertableHd tl

nodeFull :: PrimMonad m => BitPack (PrimState m) -> m Bool
nodeFull bp =
  SV.unsafeRead bp idxByte <&> (== maxNodeCapacity)

insertInPackedBS :: PrimMonad m => Word32 -> BitPack (PrimState m) -> m ()
insertInPackedBS i bp = do
  pos <- fromEnum <$> SV.unsafeRead bp idxByte
  -- the order is descending because of byteAt
  for_ [2,1,0] $ \srcWIdx ->
    let
      src = i ^. byteAt srcWIdx
      dstWIdx = 2 - srcWIdx
      dstIdx = pos + dstWIdx
    in
      SV.unsafeWrite bp dstIdx src
  SV.unsafeWrite bp idxByte (fromIntegral pos + 3)

extract :: PrimMonad m => BitPack (PrimState m) -> m (Vector Word32)
extract bp = do
  numElems <- SV.unsafeRead bp idxByte <&> flip quot 3 . fromIntegral
  V.unfoldrExactNM
    numElems
    (\n -> do
       first <- SV.unsafeRead bp (n*3) <&> fromIntegral
       mid <- SV.unsafeRead bp ((n*3)+1) <&> fromIntegral
       last <- SV.unsafeRead bp ((n*3)+2) <&> fromIntegral
       let
         i = (first .<<. 16) .|. (mid .<<. 8) .|. last
       return (i, n+1))
    0
  -- for [0..(numElems-1)] $ \n -> do
--     -- get bytes at n*3, n*3+1, n*3+2 and form an int
--     first <- SV.unsafeRead bp (n*3) <&> fromIntegral
--     mid <- SV.unsafeRead bp ((n*3)+1) <&> fromIntegral
--     last <- SV.unsafeRead bp ((n*3)+2) <&> fromIntegral
--     return $
--       (first .<<. 16) .|. (mid .<<. 8) .|. last

iteratePackedBS :: PrimMonad m => (Vector Word32 -> m ()) -> PackedList (PrimState m) -> m ()
iteratePackedBS _ Nil = return ()
iteratePackedBS action (PackedCons hd tl) = do
  iteratePackedBS action tl
  extract hd >>= action

-- iteratePackedBS' :: ([Word32] -> IO ()) -> (forall s. PackedList (PrimState (ST s))) -> IO ()
-- iteratePackedBS' _ Nil = return ()
-- iteratePackedBS' action (PackedCons hd tl) = do
--   iteratePackedBS' action tl
--   action $ runST $ extract hd

main :: IO ()
main = do
--   undefined
  pl <- foldM (flip cons) Nil [1..1000]
  iteratePackedBS print pl