content: {-# language OverloadedStrings #-}
  
  {-|
  A `SymTable` is a custom mutable hash table which is keyed by spans pointing to
  a particular `ByteString`. Every `SymTable` is initialized from a `ByteString`.
  The components of the `ByteString` are stored losslessly besides the table (but
  in a different layout).
  
  It is important that we get rid of unnecessary copying on common table
  operations. A table is defined as a nesting of mutable references. Writing the
  table size or modifying the underlying array of buckets requires no heap
  allocation.
  -}
  
  module SymTable (
      SymTable(..)
    , new
    , SymTable.lookup
    , lookupByteString
    , deleteWithHash
    , insertWithHash
    , updateWithHash
    , insert
    , delete
    , size
    , src
    , eob
    , hash
    , hashByteString
    , assocs
    , buckets
    , Entry(..)
    , loadFactor
    , loadFactor'
    ) where
  
  import qualified Data.Array.LI            as ALI
  import qualified Data.Array.LM            as ALM
  import qualified Data.ByteString          as B
  import qualified Data.ByteString.Internal as B
  import qualified Data.ByteString.Unsafe   as B
  import qualified Data.Ref.FFF             as RFFF
  import qualified Data.Ref.L               as RL
  import qualified Data.Ref.UUU             as RUUU
  
  import Data.Bits
  import Data.Word
  import GHC.Exts
  import GHC.ForeignPtr
  
  import Common
  import CoreTypes
  
  
  --------------------------------------------------------------------------------
  
  -- TODO: factor out all the shared reads from insert/delete for Cxt.Extension
  
  -- Symtable entries
  --------------------------------------------------------------------------------
  
  data Entry
    -- type, type val, def, cached (TopVar x v)
    = Top Ty {-# unpack #-} GTy Tm Tm
    -- level, type val
    | Local Lvl {-# unpack #-} GTy
  
  instance Show Entry where
    showsPrec d (Local x _)   = showParen (d > 10) (("Loc " ++ show x)++)
    showsPrec d (Top x _ _ _) = showParen (d > 10) (("Top " ++ show x)++)
  
  -- Span hashing
  --------------------------------------------------------------------------------
  
  newtype Hash = Hash {unHash :: Word}
    deriving (Eq, Show, Ord, Num, Bits) via Word
  
  hashToInt :: Hash -> Int
  hashToInt (Hash w) = fromIntegral w
  {-# inline hashToInt #-}
  
  unW# :: Word -> Word#
  unW# (W# x) = x
  {-# inline unW# #-}
  
  foldedMultiply# :: Word# -> Word# -> Word#
  foldedMultiply# x y = case timesWord2# x y of
    (# hi, lo #) -> xor# hi lo
  {-# inline foldedMultiply# #-}
  
  multiple :: Word
  multiple = 11400714819323198549
  {-# inline multiple #-}
  
  salt :: Word
  salt = 3032525626373534813
  {-# inline salt #-}
  
  combine# :: Word# -> Word# -> Word#
  combine# x y = foldedMultiply# (xor# x y) (unW# multiple)
  {-# inline combine# #-}
  
  goHash :: Word# -> Addr# -> Int# -> Word#
  goHash hash ptr len = case len <# 8# of
    1# -> case len of
      0# -> hash
      _  -> combine# hash (indexPartialWord# len ptr)
    _  -> goHash (combine# hash (indexWordOffAddr# ptr 0#)) (plusAddr# ptr 8#) (len -# 8#)
  
  goHash' :: Word# -> Addr# -> Int# -> Word#
  goHash' hash ptr len = case len <# 8# of
    1# -> case len of
      0# -> hash
      _  -> combine# hash (indexPartialWord'# len ptr)
    _  -> goHash' (combine# hash (indexWordOffAddr# ptr 0#)) (plusAddr# ptr 8#) (len -# 8#)
  
  hash# :: Addr# -> Span -> Hash
  hash# eob (Span (Pos (I# x)) (Pos (I# y))) = let
    start = plusAddr# eob (negateInt# x)
    len   = x -# y
    in case y <# 8# of
      1# -> Hash (W# (goHash' (unW# salt) start len))
      _  -> Hash (W# (goHash (unW# salt) start len))
  {-# inline hash# #-}
  
  hash :: SymTable -> Span -> IO Hash
  hash (SymTable tbl) k = do
    Ptr src  <- RFFF.readSnd =<< RUUU.readFst tbl
    pure $! hash# src k
  {-# inline hash #-}
  
  hashByteString :: Src -> IO Hash
  hashByteString str = B.unsafeUseAsCString str \(Ptr addr) -> do
    let !(I# l) = B.length str
    pure $! hash# (plusAddr# addr l) (Span (Pos (I# l)) (Pos 0))
  
  
  -- Buckets
  --------------------------------------------------------------------------------
  
  data Bucket = Empty | Cons Hash {-# unpack #-} Span Entry Bucket
  
  foldlBucket :: (b -> Hash -> Span -> Entry -> b) -> b -> Bucket -> b
  foldlBucket f acc b = go acc b where
    go acc Empty          = acc
    go acc (Cons h k v b) = let acc' = f acc h k v in go acc' b
  {-# inline foldlBucket #-}
  
  deleteFromBucket :: Addr# -> Span -> Bucket -> (# Bucket, (# Entry | (# #) #) #)
  deleteFromBucket = go where
    go src k topB = case topB of
      Empty -> (# Empty, (# | (# #) #) #)
      Cons h' k' v' b
        | 1# <- eqSpan# src k k' -> (# b, (# v' | #) #)
        | otherwise ->
           let !(# !b', del #) = go src k b
           in if ptrEq b b' then (# topB, del #) else (# Cons h' k' v' b', del #)
  
  lookupBucket :: Addr# -> Span -> Bucket -> (# Entry | (# #) #)
  lookupBucket src k = \case
    Empty -> (# | (# #) #)
    Cons h' k' v' b
      | 1# <- eqSpan# src k k' -> (# v' | #)
      | otherwise              -> lookupBucket src k b
  
  lookupBSBucket :: Addr# -> Addr# -> Span -> Bucket -> (# Entry | (# #) #)
  lookupBSBucket src' src k = \case
    Empty -> (# | (# #) #)
    Cons h' k' v' b
      | 1# <- eqBasedSpan# src k src' k' -> (# v' | #)
      | otherwise                        -> lookupBSBucket src' src k b
  
  writeBucketAtIx :: Int -> Hash -> Span -> Entry -> Bucket -> Bucket
  writeBucketAtIx i h k v e = case (,) $$! i $$! e of
    (0, Cons _ _ _ b   ) -> Cons h k v b
    (i, Cons h' k' v' b) -> Cons h' k' v' (writeBucketAtIx (i - 1) h k v b)
    (_, _              ) -> undefined
  
  insertToBucket
    :: Addr# -> Hash -> Span -> Entry -> Bucket -> (# Bucket, (# Entry | (# #) #) #)
  insertToBucket src h k ~v ~b = go src 0 h k v b b where
    go :: Addr# -> Int -> Hash -> Span -> Entry -> Bucket -> Bucket -> (# Bucket, (# Entry | (# #) #) #)
    go src i h k ~v ~topB b = case b of
      Empty                -> let c = Cons h k v topB in (# c, (# | (# #) #) #)
      Cons h' k' v' b
        | 1# <- eqSpan# src k k' -> let b = writeBucketAtIx i h k v topB in (# b, (# v' | #) #)
        | otherwise              -> go src (i + 1) h k v topB b
  
  -- SymTable
  --------------------------------------------------------------------------------
  
  newtype SymTable = SymTable
    (RUUU.Ref (RFFF.Ref Int (Ptr Word8) Int)
              (ALM.Array Bucket)
              (RL.Ref ForeignPtrContents))
  
  --------------------------------------------------------------------------------
  
  initSlotsBits :: Int
  initSlotsBits = 5
  {-# inline initSlotsBits #-}
  
  initSlots :: Int
  initSlots = unsafeShiftL 1 initSlotsBits
  {-# inline initSlots #-}
  
  eob :: SymTable -> IO (Ptr Word8)
  eob (SymTable tbl) = do
    ref <- RUUU.readFst tbl
    RFFF.readSnd ref
  {-# inline eob #-}
  
  new'# :: Int -> Ptr Word8 -> Int -> ForeignPtrContents -> IO SymTable
  new'# slots eob len fpc = do
    ref     <- RFFF.new 0 eob len
    fpcr    <- RL.new fpc
    buckets <- ALM.new slots Empty
    table   <- RUUU.new ref buckets fpcr
    pure $! SymTable table
  
  new# :: Ptr Word8 -> Int -> ForeignPtrContents -> IO SymTable
  new# = new'# initSlots
  {-# inline new# #-}
  
  new :: Src -> IO SymTable
  new (B.BS (ForeignPtr base ftc) (I# len)) =
    new# (Ptr (plusAddr# base len)) (I# len) ftc
  
  lookupByteString :: Src -> SymTable -> IO (UMaybe Entry)
  lookupByteString k (SymTable tbl) = B.unsafeUseAsCString k \(Ptr base) -> do
    let !(I# len) = B.length k
    buckets  <- RUUU.readSnd tbl
    h        <- hashByteString k
    Ptr src  <- RFFF.readSnd =<< RUUU.readFst tbl
    let bucketsSize = ALM.size buckets
        shift       = 64 - ctzInt bucketsSize
        ix          = hashToInt (unsafeShiftR h shift)
    b <- ALM.read buckets ix
    let end  = plusAddr# base len
    let span = Span (Pos (I# len)) (Pos 0)
    let !bckt = lookupBSBucket src end span b
    pure $! UMaybe# bckt
  
  lookup :: Span -> SymTable -> IO (UMaybe Entry)
  lookup k (SymTable tbl) = do
    buckets  <- RUUU.readSnd tbl
    Ptr src  <- RFFF.readSnd =<< RUUU.readFst tbl
    let bucketsSize = ALM.size buckets
        shift       = 64 - ctzInt bucketsSize
        h           = hash# src k
        ix          = hashToInt (unsafeShiftR h shift)
    b <- ALM.read buckets ix
    pure $! UMaybe# (lookupBucket src k b)
  
  resize# :: Int -> SymTable -> IO ()
  resize# bucketsSize' (SymTable tbl) = do
    buckets  <- RUUU.readSnd tbl
    buckets' <- ALM.new bucketsSize' Empty
    let shift = 64 - ctzInt bucketsSize'
    ALM.for buckets \b ->
      let go Empty          =
            pure ()
          go (Cons h k v b) = do
            ALM.modify' buckets' (hashToInt (unsafeShiftR h shift)) (Cons h k v)
            go b
      in go b
    RUUU.writeSnd tbl buckets'
  {-# noinline resize# #-}
  
  deleteWithHash :: Span -> Hash -> SymTable -> IO (UMaybe Entry)
  deleteWithHash k h (SymTable tbl) = do
    ref      <- RUUU.readFst tbl
    I# size  <- RFFF.readFst ref
    Ptr src  <- RFFF.readSnd ref
    buckets  <- RUUU.readSnd tbl
    let bucketsSize = ALM.size buckets
        shift       = 64 - ctzInt bucketsSize
        ix          = hashToInt (unsafeShiftR h shift)
    b <- ALM.read buckets ix
    let !(# !b', old #) = deleteFromBucket src k b
    let size' = I# size - (2 - tag (UMaybe# old))
    ALM.write buckets ix b'
    RFFF.writeFst ref size'
    let downsize = unsafeShiftR bucketsSize 3
    when (size' <= downsize && downsize >= initSlots) $
      resize# (unsafeShiftR bucketsSize 1) (SymTable tbl)
    pure $! UMaybe# old
  
  insertWithHash :: Span -> Hash -> Entry -> SymTable -> IO (UMaybe Entry)
  insertWithHash k h v (SymTable tbl) = do
    ref      <- RUUU.readFst tbl
    I# size  <- RFFF.readFst ref
    Ptr src  <- RFFF.readSnd ref
    buckets  <- RUUU.readSnd tbl
    let bucketsSize = ALM.size buckets
        shift       = 64 - ctzInt bucketsSize
        ix          = hashToInt (unsafeShiftR h shift)
    b <- ALM.read buckets ix
    let !(# b', old #) = insertToBucket src h k v b
    let size' = I# size + tag (UMaybe# old) - 1
    ALM.write buckets ix b'
    RFFF.writeFst ref size'
    when (size' >= unsafeShiftR bucketsSize 1) $
      resize# (unsafeShiftL bucketsSize 1) (SymTable tbl)
    pure $! UMaybe# old
  
  updateWithHash :: Span -> Hash -> UMaybe Entry -> SymTable -> IO (UMaybe Entry)
  updateWithHash k h mv tbl = case mv of
    UNothing -> do
      -- debug ["deletewithhash", showSpan (src tbl) k]
      deleteWithHash k h tbl
    UJust v  -> do
      -- debug ["insertwithhash", showSpan (src tbl) k]
      insertWithHash k h v tbl
  {-# inline updateWithHash #-}
  
  insert :: Span -> Entry -> SymTable -> IO (UMaybe Entry)
  insert k v tbl = do
    h <- hash tbl k
    insertWithHash k h v tbl
  
  delete :: Span -> SymTable -> IO (UMaybe Entry)
  delete k tbl = do
    h <- hash tbl k
    deleteWithHash k h tbl
  
  size :: SymTable -> IO Int
  size (SymTable tbl) = RFFF.readFst =<< RUUU.readFst tbl
  {-# inline size #-}
  
  src :: SymTable -> Src
  src (SymTable tbl) = runIO do
    ref     <- RUUU.readFst tbl
    Ptr end <- RFFF.readSnd ref
    I# len  <- RFFF.readThd ref
    fptr    <- RL.read =<< RUUU.readThd tbl
    let start = plusAddr# end (negateInt# len)
    pure $! B.BS (ForeignPtr start fptr) (I# len)
  
  -- testing
  --------------------------------------------------------------------------------
  
  loadFactor :: SymTable -> IO Double
  loadFactor (SymTable tbl) = do
    ref      <- RUUU.readFst tbl
    size     <- RFFF.readFst ref
    buckets  <- RUUU.readSnd tbl
    let bucketsSize = ALM.size buckets
    pure $! fromIntegral size / fromIntegral bucketsSize
  
  loadFactor' :: SymTable -> IO Double
  loadFactor' tbl = do
    bs <- buckets tbl
    let blen = length bs
    let size = length $ concat bs
    pure $! fromIntegral size / fromIntegral blen
  
  assocs :: SymTable -> IO [(String, Entry)]
  assocs stbl@(SymTable tbl) = do
    buckets <- ALM.freeze =<< RUUU.readSnd tbl
    pure $! ALI.foldl'
      (\acc b -> foldlBucket
        (\acc h k v -> (showSpan (src stbl) k, v):acc) acc b)
        [] buckets
  
  buckets :: SymTable -> IO [[(Hash, String, Entry)]]
  buckets stbl@(SymTable tbl) = do
    buckets <- ALM.freeze =<< RUUU.readSnd tbl
    pure $! ALI.foldl'
      (\acc b -> foldlBucket
          (\acc h k v -> (h, showSpan (src stbl) k, v):acc) [] b : acc)
          [] buckets
  
  testHash :: Src -> Span -> Hash
  testHash str s = runIO $ B.unsafeUseAsCString str \(Ptr addr) -> do
    let !(I# l) = B.length str
        eob = plusAddr# addr l
    pure $! hash# eob s
  
  testEqSpan :: Src -> Span -> Span -> Bool
  testEqSpan str s s' = runIO $ B.unsafeUseAsCString str \(Ptr addr) -> do
    let !(I# l) = B.length str
        eob = plusAddr# addr l
    pure $! isTrue# (eqSpan# eob s s')
  
  -- test = do
  --   tbl <- new "EvalCon0EvalCon0        "
  --   insert (Span (Pos 24) (Pos 16)) (Local 10 (gjoin VU)) tbl
  --   insert (Span (Pos 16) (Pos 8)) (Local 10 (gjoin VU)) tbl
  --   -- insert (Span (Pos 8) (Pos 0)) (Local 10 (gjoin VU)) tbl
  --   -- lk <- SymTable.lookup (Span (Pos 16) (Pos 8)) tbl
  --   -- print lk
  
  --   mapM_ print =<< buckets tbl
  