wtf.hs

module Main where

import System.TimeIt
import Data.ByteString.Lazy qualified as LBS
import Data.ByteString qualified as BS
import Data.ByteString.Lazy (ByteString)
import Crypto.Random (getRandomBytes)
import Control.Monad
import Data.Functor
import Data.Function
import UnliftIO
import Data.List (unfoldr)
import Data.Foldable
import Data.Traversable
import System.FilePath
import Data.Map qualified as Map
import Data.Map (Map)
import Data.Hashable
import Data.Vector qualified as V
import Data.Vector ((!))
import Safe
import Control.Concurrent.STM (flushTQueue)
-- import Data.IORef (atomicallyModifyIORef
import Prettyprinter
import Data.HashMap.Strict (HashMap)
import Data.HashMap.Strict qualified as HashMap
import System.Random.Stateful
-- .MWC qualified as MWC
import System.Random.MWC qualified as MWC

import StmContainers.Map qualified as StmMap
import Control.Concurrent.STM.TSkipList as TS
import Data.Trie as Trie

import DBPipe.SQLite as SQL

import Database.RocksDB qualified as R

import Data.CritBit.Map.Lazy qualified as C
import Data.Word
import Data.Int
import Data.Default
import Codec.Compression.GZip as GZip

import System.Directory

-- data SkipListNode =
--   SkipListNode
--   {
--   }

runUpdates :: MonadUnliftIO m => Int -> Int -> Int -> (ByteString -> m ()) -> m ()
runUpdates total k n fn = do
  tot <- newTVarIO total

  let (a,b) = total `divMod` k

  let tasks = a + b : replicate (k-1) a

  g <- liftIO $ MWC.createSystemRandom

  threads <- for tasks $ \t -> async do
               replicateM_ t do
                   bs <- liftIO $ replicateM n (MWC.uniformM @Word8 g) <&> LBS.pack
                   fn bs

  mapM_ wait threads

randomRead :: (StatefulGen g IO) => Int -> g -> Double -> (ByteString -> IO (Maybe ByteString)) -> IO (Maybe ByteString)
randomRead n g p rfn = do
  k <- getRandomBytes n <&> LBS.fromStrict
  pr <- MWC.uniformRM (0, 1.0 :: Double)  g

  if pr > (1 - p) then do
    rfn k
  else
    pure Nothing

main :: IO ()
main = do

  let threads = 8
  -- let total = 1_000_000
  let total = 1_000_000
  let keylen = 32

  g <- MWC.createSystemRandom

  liftIO $ removePathForcibly "lsm.wal"
  liftIO $ removePathForcibly "lsm.data"
  liftIO $ removePathForcibly "jopakita.db"
  liftIO $ removePathForcibly "jopakita"

  timeItNamed "just-none" do
    runUpdates total threads keylen $ \_ -> do
      pure ()

  timeItNamed "just-count" do
    tm <- newIORef 0
    runUpdates total threads keylen $ \_ -> do
      atomicModifyIORef' tm (\x -> (succ x, ()))
    readIORef tm >>= print

  timeItNamed "just-tvar-map" do
    tvm <- newTVarIO mempty

    runUpdates total threads keylen $ \bs -> do
      atomically $ modifyTVar' tvm (Map.insert bs bs)

  timeItNamed "just-ioref-map" do
    tvm <- newIORef mempty
    runUpdates total threads keylen $ \bs -> do
      atomicModifyIORef' tvm (\x -> (Map.insert bs bs x, ()))

  timeItNamed "just-tvar-hashmap" do
    tvm <- newTVarIO mempty
    runUpdates total threads keylen $ \bs -> do
      atomically $ modifyTVar' tvm (HashMap.insert bs bs)

  timeItNamed "just-tvar-buck-map" do
    let buc = threads
    tvm <- replicateM buc (newTVarIO mempty) <&> V.fromList
    runUpdates total threads keylen $ \bs -> do
      let i = maybe 0 fst (LBS.uncons bs) `mod` fromIntegral buc
      let t = tvm ! fromIntegral i
      atomically $ modifyTVar' t (Map.insert bs bs)


  timeItNamed "just-tvar-buck-map-disk-wal" do

    let buc = threads
    tvm <- replicateM buc (newTVarIO mempty) <&> V.fromList

    walq <- newTQueueIO

    walp <- async do
              tw <- newIORef ( 0 :: Int64 )

              withBinaryFile "lsm.wal" AppendMode $ \ha -> do
                fix \next -> do
                  it <- atomically $ readTQueue walq
                  case it of
                    Nothing -> pure ()
                    Just (k,v) -> do
                      let ss = k <> v
                      LBS.hPutStr ha ss

                      modifyIORef tw (+ LBS.length ss)
                      sz <- readIORef tw
                      -- sz <- atomicModifyIORef tw (\x -> (x + LBS.length ss, x + LBS.length ss))

                      when (sz `mod` (1024*1024) == 0) do
                        hFlush ha

                      next

    runUpdates total threads keylen $ \bs -> do
      let i = maybe 0 fst (LBS.uncons bs) `mod` fromIntegral buc
      let t = tvm ! fromIntegral i
      atomically $ modifyTVar' t (Map.insert bs bs)
      atomically $ writeTQueue walq (Just (bs,bs))

    withBinaryFile "lsm.data" AppendMode $ \ha -> do
      joined  <- mapM readTVarIO tvm <&> Map.toList . Map.unions
      for_ joined $ \(k,v) -> do
          LBS.hPutStr ha (k <> v)

    atomically $ writeTQueue walq Nothing

    wait walp


  timeItNamed "just-stm-map" do
    tm <- StmMap.newIO @ByteString @ByteString
    runUpdates total 8 keylen $ \bs -> do
      atomically $ StmMap.insert bs bs tm

  timeItNamed "bs-trie" do
    tm <- newTVarIO Trie.empty
    runUpdates total threads keylen $ \bs -> do
      let k = LBS.toStrict bs
      atomically $ modifyTVar' tm (Trie.insert k bs)

  timeItNamed "critbit" do
    tm <- newTVarIO C.empty
    runUpdates total threads keylen $ \bs -> do
      let k = LBS.toStrict bs
      atomically $ modifyTVar' tm (C.insert k bs)

  timeItNamed "qwriter" do
    -- tm <- newTVarIO mempty
    q  <- newTQueueIO
    z <- newIORef mempty

    runUpdates total threads keylen $ \bs -> do
      let k = LBS.toStrict bs
      atomicModifyIORef z (\x -> (k : x, ()))
      -- atomically $ writeTQueue q k -- modifyTVar tm (C.insert k bs)

  env <- newDBPipeEnv dbPipeOptsDef "jopakita.db"

  withDB env do
    ddl "create table kv (k text not null, v text not null, primary key (k))"
    commitAll

  timeItNamed "sqlite" do
    withDB env $ transactional do
      runUpdates total threads keylen $ \bs -> do
        withDB env $ do
          SQL.insert "insert into kv (k,v) values (?,?) on conflict (k) do nothing" (bs,bs)
      withDB env commitAll

    R.withDB "jopakita" (def { R.createIfMissing = True }) $ \db -> do
      timeItNamed "rocks" do
        R.withSnapshot db $ \db1 -> do
          runUpdates total threads keylen $ \bs -> do
            R.put db1 (LBS.toStrict bs) (LBS.toStrict bs)
            pure ()

  pure ()

  -- timeItNamed "just-tskiplist 8" do
  --   tm <- TS.newIO @ByteString @ByteString
  --   runUpdates total 8 keylen $ \bs -> do
  --     atomically $ TS.insert bs bs tm