CAttrSet.hs

-- | C-backed sorted attribute set via FFI.
--
-- Wraps @cbits/nn_attrset.c@ - a contiguous sorted array of
-- @(nn_symbol_t, void*)@ pairs.  Replaces Haskell's @Data.Map.Strict@
-- for attribute sets, cutting per-entry overhead from ~48 bytes
-- (Map.Bin node) to ~12 bytes (4-byte symbol + 8-byte pointer).
--
-- Construction is two-phase: insert entries, then freeze (sort + dedup).
-- After freeze, lookup is O(log n) binary search over contiguous memory.
--
-- Values are opaque @CThunkPtr@ handles stored as @void*@ on the C side.
-- The C side never dereferences them.  Thunk lifetimes are managed by
-- the thunk arena; attribute set cleanup only frees the key\/value arrays.
--
-- A set holds at most 'maxBound' :: 'Word32' entries, upstream's own
-- ceiling (its @Bindings@ counts in @uint32_t@).  Every constructor and
-- insert checks its C result, so a set the C side cannot build raises
-- 'Nix.Eval.CStatus.CStatusError' instead of leaving a NULL or a short
-- set behind.
module Nix.Eval.CAttrSet
  ( -- * Opaque handle
    CAttrSet,

    -- * Lifecycle
    cattrsetNew,

    -- * Construction
    cattrsetInsert,
    cattrsetFreeze,

    -- * Query
    cattrsetLookup,
    cattrsetIndex,
    cattrsetGetValue,
    cattrsetSetValue,
    cattrsetGetKey,

    -- * Bulk access
    cattrsetSize,
    cattrsetKeys,

    -- * Set operations
    cattrsetUnion,
    cattrsetRemoveKeys,
  )
where

import Control.Monad (unless)
import Data.Word (Word32)
import Foreign.C.Types (CInt (..))
import Foreign.Marshal.Array (peekArray, withArray)
import Foreign.Ptr (Ptr, castPtr, nullPtr)
import Nix.Eval.CStatus (cStatusFailure, checkedCPtr)
import Nix.Eval.CThunk (CThunkPtr)
import Nix.Eval.Symbol (Symbol (..))

-- | Opaque handle to a C-allocated attribute set.
-- The Haskell side holds a raw pointer; the C side owns the memory.
data NnAttrSet

-- | Convenience alias.
type CAttrSet = Ptr NnAttrSet

-- ---------------------------------------------------------------------------
-- FFI imports (all unsafe - no callbacks, fast data access)
-- ---------------------------------------------------------------------------

foreign import ccall unsafe "nn_attrset_new"
  c_nn_attrset_new :: Word32 -> IO CAttrSet

foreign import ccall unsafe "nn_attrset_insert"
  c_nn_attrset_insert :: CAttrSet -> Word32 -> Ptr () -> IO CInt

foreign import ccall unsafe "nn_attrset_freeze"
  c_nn_attrset_freeze :: CAttrSet -> IO ()

foreign import ccall unsafe "nn_attrset_lookup"
  c_nn_attrset_lookup :: CAttrSet -> Word32 -> IO (Ptr ())

foreign import ccall unsafe "nn_attrset_index"
  c_nn_attrset_index :: CAttrSet -> Word32 -> IO Word32

foreign import ccall unsafe "nn_attrset_set_value"
  c_nn_attrset_set_value :: CAttrSet -> Word32 -> Ptr () -> IO ()

foreign import ccall unsafe "nn_attrset_get_value"
  c_nn_attrset_get_value :: CAttrSet -> Word32 -> IO (Ptr ())

foreign import ccall unsafe "nn_attrset_get_key"
  c_nn_attrset_get_key :: CAttrSet -> Word32 -> IO Word32

foreign import ccall unsafe "nn_attrset_size"
  c_nn_attrset_size :: CAttrSet -> IO Word32

foreign import ccall unsafe "nn_attrset_keys_ptr"
  c_nn_attrset_keys_ptr :: CAttrSet -> IO (Ptr Word32)

foreign import ccall unsafe "nn_attrset_union"
  c_nn_attrset_union :: CAttrSet -> CAttrSet -> IO CAttrSet

foreign import ccall unsafe "nn_attrset_remove_keys"
  c_nn_attrset_remove_keys :: CAttrSet -> Ptr Word32 -> Word32 -> IO CAttrSet

-- | @NN_ATTRSET_NOT_FOUND@, the index answer for an absent key.  Must
-- stay in lockstep with @cbits\/nn_attrset.h@.
indexNotFound :: Word32
indexNotFound = maxBound

-- ---------------------------------------------------------------------------
-- Lifecycle
-- ---------------------------------------------------------------------------

-- | Allocate a new empty attribute set sized for the expected number of
-- entries.  The count only sizes the first allocation, so one past the
-- ceiling is clamped to it: the inserts are what enforce the ceiling.
cattrsetNew :: Int -> IO CAttrSet
cattrsetNew expected = checkedCPtr "nn_attrset_new" =<< c_nn_attrset_new capacity
  where
    capacity = fromIntegral (max 0 (min expected entryCeiling))
    entryCeiling = fromIntegral (maxBound :: Word32)

-- ---------------------------------------------------------------------------
-- Construction
-- ---------------------------------------------------------------------------

-- | Insert a key-value pair.  Call before freeze.
-- The value is a CThunkPtr cast to Ptr () - C never dereferences it.
cattrsetInsert :: CAttrSet -> Symbol -> CThunkPtr -> IO ()
cattrsetInsert set (Symbol sym) ptr = do
  status <- c_nn_attrset_insert set sym (castPtr ptr)
  unless (status == 0) $
    cStatusFailure "nn_attrset_insert" "the attribute set cannot grow (allocation failure or its 4294967295-entry ceiling)"

-- | Sort and deduplicate.  Must be called once before any queries.
cattrsetFreeze :: CAttrSet -> IO ()
cattrsetFreeze = c_nn_attrset_freeze

-- ---------------------------------------------------------------------------
-- Query
-- ---------------------------------------------------------------------------

-- | Look up a key.  Returns 'Nothing' if not found.
-- The returned 'CThunkPtr' is borrowed - the thunk arena owns it.
cattrsetLookup :: CAttrSet -> Symbol -> IO (Maybe CThunkPtr)
cattrsetLookup set (Symbol sym) = do
  ptr <- c_nn_attrset_lookup set sym
  pure (if ptr == nullPtr then Nothing else Just (castPtr ptr))

-- | Find the index of a key, or 'Nothing' if not found.
cattrsetIndex :: CAttrSet -> Symbol -> IO (Maybe Word32)
cattrsetIndex set (Symbol sym) = do
  idx <- c_nn_attrset_index set sym
  pure (if idx == indexNotFound then Nothing else Just idx)

-- | Get the value at a known index.
cattrsetGetValue :: CAttrSet -> Word32 -> IO CThunkPtr
cattrsetGetValue set idx = castPtr <$> c_nn_attrset_get_value set idx

-- | Update the value at a known index (for two-phase construction).
cattrsetSetValue :: CAttrSet -> Word32 -> CThunkPtr -> IO ()
cattrsetSetValue set idx ptr = c_nn_attrset_set_value set idx (castPtr ptr)

-- | Get the key symbol at a known index.
cattrsetGetKey :: CAttrSet -> Word32 -> IO Symbol
cattrsetGetKey set idx = Symbol <$> c_nn_attrset_get_key set idx

-- ---------------------------------------------------------------------------
-- Bulk access
-- ---------------------------------------------------------------------------

-- | Number of entries (unique keys after freeze).
cattrsetSize :: CAttrSet -> IO Word32
cattrsetSize = c_nn_attrset_size

-- | Read all keys as a list of symbols (sorted).
cattrsetKeys :: CAttrSet -> IO [Symbol]
cattrsetKeys set = do
  n <- cattrsetSize set
  ptr <- c_nn_attrset_keys_ptr set
  raw <- peekArray (fromIntegral n) ptr
  pure (map Symbol raw)

-- ---------------------------------------------------------------------------
-- Set operations
-- ---------------------------------------------------------------------------

-- | Right-biased union (// semantics).  Result is frozen.
cattrsetUnion :: CAttrSet -> CAttrSet -> IO CAttrSet
cattrsetUnion a b = checkedCPtr "nn_attrset_union" =<< c_nn_attrset_union a b

-- | Remove a list of keys.  Result is frozen.
cattrsetRemoveKeys :: CAttrSet -> [Symbol] -> IO CAttrSet
cattrsetRemoveKeys set syms = do
  let raw = map unSymbol syms
      n = fromIntegral (length raw)
  withArray raw $ \ptr ->
    checkedCPtr "nn_attrset_remove_keys" =<< c_nn_attrset_remove_keys set ptr n