LMDB.hs

{-# OPTIONS_HADDOCK prune #-}
{-# LANGUAGE BangPatterns   #-}
{-# LANGUAGE EmptyDataDecls #-}

-- |
-- Module: Database.LMDB
-- Copyright: (c) 2026 Jared Tobin
-- License: MIT
-- Maintainer: Jared Tobin <jared@ppad.tech>
--
-- Minimal bindings to
-- [LMDB](https://www.symas.com/lmdb), an embedded ACID key-value
-- store. The upstream LMDB C source is vendored under @cbits\/@ at
-- release @LMDB_0.9.33@; no external @liblmdb@ is needed.
--
-- LMDB has a single-writer, many-reader transaction model and uses
-- memory mapping for zero-copy reads. This module presents that as a
-- bracketed @IO@-only interface.
--
-- Read-only and read-write transactions are distinguished at the type
-- level via a phantom parameter on 'Txn', so using a write primitive
-- (e.g. 'put') on a read transaction is a compile-time error.
--
-- Most operations throw 'LMDBException' on error. 'get' and 'del' are
-- the exceptions: they map a missing key onto 'Nothing' and 'False'
-- respectively, since absence is normal control flow.

module Database.LMDB (
  -- * Environments
    Env
  , EnvFlags(..)
  , defaultEnvFlags
  , open
  , close
  , withEnv

  -- * Transactions
  , Txn
  , RO
  , RW
  , withReadTxn
  , withWriteTxn

  -- * Databases
  , Dbi
  , openDbi

  -- * Key-value operations
  , get
  , put
  , del

  -- * Cursors
  , Cursor
  , withCursor
  , cursorFirst
  , cursorLast
  , cursorNext
  , cursorPrev
  , cursorSeek

  -- * Errors
  , LMDBException(..)
  ) where

import Control.Exception
  ( Exception, bracket, bracketOnError, mask, onException, throwIO)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Unsafe as BU
import Foreign.C.String (peekCString, withCString)
import Foreign.C.Types (CInt, CSize, CUInt)
import Foreign.Marshal.Alloc (alloca, allocaBytes)
import Foreign.Ptr (Ptr, nullPtr, plusPtr)
import Foreign.Storable (peek, peekByteOff, pokeByteOff)
import Database.LMDB.Internal

-- handles --------------------------------------------------------------------

-- | An LMDB environment. Roughly, an open database file (or, with
--   'envNoSubdir' off, a directory containing one).
newtype Env = Env (Ptr MDB_env)

-- | Phantom tag for read-only transactions.
data RO

-- | Phantom tag for read-write transactions.
data RW

-- | A transaction. The phantom @s@ distinguishes 'RO' from 'RW' so
--   that write operations refuse to type-check against read-only
--   transactions.
newtype Txn s = Txn (Ptr MDB_txn)

-- | A database identifier within an 'Env'.
newtype Dbi = Dbi MDB_dbi

-- | A cursor for range-scanning a 'Dbi'.
newtype Cursor s = Cursor (Ptr MDB_cursor)

-- environment flags ----------------------------------------------------------

-- | Configuration for 'open'.
data EnvFlags = EnvFlags
  { envMapSize  :: !Int
    -- ^ map size in bytes; LMDB will refuse writes beyond this. The
    --   default of 10 MiB is intentionally small — set it to something
    --   appropriate for your workload before opening.
  , envMaxDbs   :: !Int
    -- ^ maximum number of named sub-databases. Default 1.
  , envNoSubdir :: !Bool
    -- ^ if 'True', treat the path as a regular file rather than a
    --   directory containing @data.mdb@ and @lock.mdb@.
  , envReadOnly :: !Bool
    -- ^ open the environment read-only.
  , envNoSync   :: !Bool
    -- ^ skip @fsync@ at commit time. Faster, but a crash may lose
    --   the last transaction. The database remains internally
    --   consistent.
  } deriving Show

-- | Default 'EnvFlags': 10 MiB map, one database, directory layout,
--   read-write, syncing at commit.
--
--   >>> envMapSize defaultEnvFlags
--   10485760
defaultEnvFlags :: EnvFlags
defaultEnvFlags = EnvFlags
  { envMapSize  = 10 * 1024 * 1024
  , envMaxDbs   = 1
  , envNoSubdir = False
  , envReadOnly = False
  , envNoSync   = False
  }

envFlagBits :: EnvFlags -> CUInt
envFlagBits EnvFlags { envNoSubdir = nsd, envReadOnly = ro, envNoSync = ns } =
  (if nsd then _MDB_NOSUBDIR else 0)
    + (if ro then _MDB_RDONLY   else 0)
    + (if ns then _MDB_NOSYNC   else 0)

-- error handling -------------------------------------------------------------

-- | A typed wrapper around LMDB error codes. 'LMDBOther' carries the
--   raw error code plus the message from @mdb_strerror@.
data LMDBException
  = LMDBKeyExist
  | LMDBNotFound          -- ^ Only seen via "Database.LMDB.Internal";
                          --   the safe layer maps this to 'Nothing'
                          --   or 'False' instead of throwing.
  | LMDBMapFull
  | LMDBCorrupted
  | LMDBPanic
  | LMDBVersionMismatch
  | LMDBOther !Int !String
  deriving Show

instance Exception LMDBException

throwLmdb :: CInt -> IO a
throwLmdb code
  | code == _MDB_KEYEXIST         = throwIO LMDBKeyExist
  | code == _MDB_NOTFOUND         = throwIO LMDBNotFound
  | code == _MDB_MAP_FULL         = throwIO LMDBMapFull
  | code == _MDB_CORRUPTED        = throwIO LMDBCorrupted
  | code == _MDB_PANIC            = throwIO LMDBPanic
  | code == _MDB_VERSION_MISMATCH = throwIO LMDBVersionMismatch
  | otherwise = do
      msg <- mdb_strerror code >>= peekCString
      throwIO (LMDBOther (fromIntegral code) msg)

check :: CInt -> IO ()
check 0 = pure ()
check c = throwLmdb c
{-# INLINE check #-}

-- environment lifecycle ------------------------------------------------------

-- | Open an LMDB environment at the given path. The path must already
--   exist (as a directory unless 'envNoSubdir' is set, in which case
--   the file's parent directory must exist).
--
--   Throws 'LMDBException' on failure.
--
--   >>> env <- open "/tmp/mydb" defaultEnvFlags { envNoSubdir = True }
open :: FilePath -> EnvFlags -> IO Env
open path flags = mask $ \restore -> do
  envp <- alloca $ \pp -> do
    check =<< mdb_env_create pp
    peek pp
  let env = Env envp
  flip onException (close env) . restore $ do
    check =<< mdb_env_set_mapsize envp
                  (fromIntegral (envMapSize flags))
    check =<< mdb_env_set_maxdbs envp
                  (fromIntegral (envMaxDbs flags))
    withCString path $ \cpath ->
      check =<< mdb_env_open envp cpath (envFlagBits flags) 0o644
  pure env

-- | Close an environment. Must not be called while transactions are
--   live.
close :: Env -> IO ()
close (Env envp) = mdb_env_close envp

-- | Bracketed environment lifecycle: 'open' the environment, run the
--   action, 'close' on exit (including async exceptions).
--
--   >>> withEnv "/tmp/mydb" defaultEnvFlags $ \env -> ...
withEnv :: FilePath -> EnvFlags -> (Env -> IO a) -> IO a
withEnv path flags = bracket (open path flags) close

-- transactions ---------------------------------------------------------------

beginTxn :: Env -> CUInt -> IO (Ptr MDB_txn)
beginTxn (Env envp) flags = alloca $ \pp -> do
  check =<< mdb_txn_begin envp nullPtr flags pp
  peek pp

-- | Run an action in a read-only transaction. The transaction is
--   aborted on exit; read transactions have no work to commit.
--
--   >>> withReadTxn env $ \txn -> ...
withReadTxn :: Env -> (Txn RO -> IO a) -> IO a
withReadTxn env act = bracketOnError
  (beginTxn env _MDB_RDONLY)
  mdb_txn_abort
  (\txnp -> do
      r <- act (Txn txnp)
      mdb_txn_abort txnp
      pure r)

-- | Run an action in a read-write transaction. Commits if the action
--   returns normally; aborts on exception.
--
--   >>> withWriteTxn env $ \txn -> ...
withWriteTxn :: Env -> (Txn RW -> IO a) -> IO a
withWriteTxn env act = bracketOnError
  (beginTxn env 0)
  mdb_txn_abort
  (\txnp -> do
      r <- act (Txn txnp)
      check =<< mdb_txn_commit txnp
      pure r)

-- databases ------------------------------------------------------------------

-- | Open a database within a transaction. Pass 'Nothing' for the
--   unnamed default database. If the @create@ flag is 'True', the
--   database is created if it does not already exist — and the
--   transaction must be a read-write transaction.
--
--   >>> dbi <- openDbi txn Nothing True
openDbi
  :: Txn s
  -> Maybe BS.ByteString -- ^ database name; 'Nothing' for the default
  -> Bool                -- ^ create if not present
  -> IO Dbi
openDbi (Txn txnp) mname create = alloca $ \pp -> do
  let flags = if create then _MDB_CREATE else 0
      go cstr = do
        check =<< mdb_dbi_open txnp cstr flags pp
        Dbi <$> peek pp
  case mname of
    Nothing -> go nullPtr
    Just nm -> BU.unsafeUseAsCString nm go

-- bytestring <-> MDB_val helpers --------------------------------------------

-- An MDB_val is two words: a CSize length at offset 0, a payload Ptr
-- at offset 8 (on 64-bit). We poke fields directly rather than going
-- through the Storable instance so hot paths don't depend on it being
-- inlined.

withBSAsVal :: BS.ByteString -> (Ptr MDB_val -> IO a) -> IO a
withBSAsVal bs k =
  BU.unsafeUseAsCStringLen bs $ \(p, n) ->
  allocaBytes 16 $ \vp -> do
    pokeByteOff vp 0 (fromIntegral n :: CSize)
    pokeByteOff vp 8 p
    k vp
{-# INLINE withBSAsVal #-}

peekVal :: Ptr MDB_val -> IO BS.ByteString
peekVal vp = do
  sz <- peekByteOff vp 0 :: IO CSize
  pd <- peekByteOff vp 8
  BS.packCStringLen (pd, fromIntegral sz)
{-# INLINE peekVal #-}

-- key-value ops --------------------------------------------------------------

-- | Look up a key. Returns 'Nothing' if not found. Other LMDB errors
--   throw 'LMDBException'.
--
--   The returned 'BS.ByteString' is a copy of the data; LMDB's
--   memory-mapped buffer is not aliased into Haskell.
--
--   >>> get txn dbi "hello"
--   Just "world"
get
  :: Txn s
  -> Dbi
  -> BS.ByteString               -- ^ key
  -> IO (Maybe BS.ByteString)    -- ^ value, if present
get (Txn txnp) (Dbi dbi) k =
  BU.unsafeUseAsCStringLen k $ \(kdata, klen) ->
  allocaBytes 32 $ \kp -> do
    let !vp = kp `plusPtr` 16
    pokeByteOff kp 0 (fromIntegral klen :: CSize)
    pokeByteOff kp 8 kdata
    code <- mdb_get txnp dbi kp vp
    if   code == _MDB_NOTFOUND
    then pure Nothing
    else do
      check code
      Just <$> peekVal vp
{-# INLINE get #-}

-- | Insert or replace a key. Throws 'LMDBException' on failure.
--
--   >>> put txn dbi "hello" "world"
put
  :: Txn RW
  -> Dbi
  -> BS.ByteString -- ^ key
  -> BS.ByteString -- ^ value
  -> IO ()
put (Txn txnp) (Dbi dbi) k v =
  BU.unsafeUseAsCStringLen k $ \(kdata, klen) ->
  BU.unsafeUseAsCStringLen v $ \(vdata, vlen) ->
  allocaBytes 32 $ \kp -> do
    let !vp = kp `plusPtr` 16
    pokeByteOff kp 0 (fromIntegral klen :: CSize)
    pokeByteOff kp 8 kdata
    pokeByteOff vp 0 (fromIntegral vlen :: CSize)
    pokeByteOff vp 8 vdata
    check =<< mdb_put txnp dbi kp vp 0
{-# INLINE put #-}

-- | Delete a key. Returns 'True' if the key existed and was removed,
--   'False' if it was already absent.
--
--   >>> del txn dbi "hello"
--   True
del
  :: Txn RW
  -> Dbi
  -> BS.ByteString -- ^ key
  -> IO Bool       -- ^ 'True' if the key was present and removed
del (Txn txnp) (Dbi dbi) k =
  withBSAsVal k $ \kp -> do
    code <- mdb_del txnp dbi kp nullPtr
    if   code == _MDB_NOTFOUND then pure False
    else do
      check code
      pure True
{-# INLINE del #-}

-- cursors --------------------------------------------------------------------

-- | Bracketed cursor lifecycle: open a cursor over the given 'Dbi',
--   run the action, close the cursor on exit (including async
--   exceptions).
--
--   >>> withCursor txn dbi $ \cur -> ...
withCursor :: Txn s -> Dbi -> (Cursor s -> IO a) -> IO a
withCursor (Txn txnp) (Dbi dbi) act = bracket open' mdb_cursor_close use'
  where
    open' = alloca $ \pp -> do
      check =<< mdb_cursor_open txnp dbi pp
      peek pp
    use' p = act (Cursor p)

-- Scan-position helper (no input key): one 32-byte buffer for both
-- output MDB_vals.
cursorMove
  :: Cursor s
  -> CInt              -- MDB_cursor_op
  -> IO (Maybe (BS.ByteString, BS.ByteString))
cursorMove (Cursor cp) op =
  allocaBytes 32 $ \kp -> do
    let !vp = kp `plusPtr` 16
    code <- mdb_cursor_get cp kp vp op
    if   code == _MDB_NOTFOUND
    then pure Nothing
    else do
      check code
      k <- peekVal kp
      v <- peekVal vp
      pure (Just (k, v))
{-# INLINE cursorMove #-}

-- | Position at the first key/value pair, or 'Nothing' on an empty
--   database.
--
--   >>> cursorFirst cur
--   Just ("a","alpha")
cursorFirst :: Cursor s -> IO (Maybe (BS.ByteString, BS.ByteString))
cursorFirst c = cursorMove c _MDB_FIRST
{-# INLINE cursorFirst #-}

-- | Position at the last key/value pair, or 'Nothing' on an empty
--   database.
cursorLast :: Cursor s -> IO (Maybe (BS.ByteString, BS.ByteString))
cursorLast c = cursorMove c _MDB_LAST
{-# INLINE cursorLast #-}

-- | Advance to the next key/value pair, or 'Nothing' if already at
--   the end.
cursorNext :: Cursor s -> IO (Maybe (BS.ByteString, BS.ByteString))
cursorNext c = cursorMove c _MDB_NEXT
{-# INLINE cursorNext #-}

-- | Step back to the previous key/value pair, or 'Nothing' if already
--   at the start.
cursorPrev :: Cursor s -> IO (Maybe (BS.ByteString, BS.ByteString))
cursorPrev c = cursorMove c _MDB_PREV
{-# INLINE cursorPrev #-}

-- | Position at the first key /greater than or equal to/ the given
--   key, or 'Nothing' if no such key exists.
--
--   >>> cursorSeek cur "m"
--   Just ("n","november")
cursorSeek
  :: Cursor s
  -> BS.ByteString
  -> IO (Maybe (BS.ByteString, BS.ByteString))
cursorSeek (Cursor cp) k =
  BU.unsafeUseAsCStringLen k $ \(kdata, klen) ->
  allocaBytes 32 $ \kp -> do
    let !vp = kp `plusPtr` 16
    pokeByteOff kp 0 (fromIntegral klen :: CSize)
    pokeByteOff kp 8 kdata
    code <- mdb_cursor_get cp kp vp _MDB_SET_RANGE
    if   code == _MDB_NOTFOUND
    then pure Nothing
    else do
      check code
      ok <- peekVal kp
      ov <- peekVal vp
      pure (Just (ok, ov))