module Database.Muesli.Allocator
  ( empty
  , add
  , build
  , buildExtra
  , alloc
  ) where
import           Control.Exception             (throw)
import qualified Data.IntMap.Strict            as IntMap
import           Data.List                     (foldl', sortOn)
import qualified Data.Map.Strict               as Map
import           Data.Maybe                    (fromMaybe)
import           Database.Muesli.Backend.Types (LogRecord (..))
import           Database.Muesli.State         (GapsIndex, MainIndex)
import           Database.Muesli.Types         (DatabaseError (..), DocAddress,
                                                DocSize)
empty :: DocAddress -> GapsIndex
empty addr = Map.singleton (maxBound  addr) [addr]
add :: DocSize -> DocAddress -> GapsIndex -> GapsIndex
add sz addr gs = Map.insert sz (addr:as) gs
  where as = fromMaybe [] $ Map.lookup sz gs
build :: MainIndex -> GapsIndex
build idx = addTail . foldl' f (Map.empty, 0) . sortOn recAddress .
            filter (not . recDeleted) . map head $ IntMap.elems idx
  where
    f (gs, addr) r = (gs', recAddress r + recSize r)
      where gs' = if addr == recAddress r then gs
                  else add sz addr gs
            sz = recAddress r  addr
    addTail (gs, addr) = add (maxBound  addr) addr gs
buildExtra :: DocAddress -> [LogRecord] -> GapsIndex
buildExtra pos = foldl' f (empty pos)
  where f gs r = add (recSize r) (recAddress r) gs
alloc :: GapsIndex -> DocSize -> (DocAddress, GapsIndex)
alloc gs sz =
  case Map.lookupGE sz gs of
    Nothing -> throw $ DataAllocationError sz (fst <$> Map.lookupLT maxBound gs)
                       "Data allocation error."
    Just (gsz, a:as) ->
      if delta == 0 then (a, gs')
      else (a, add delta (a + sz) gs')
      where gs' = if null as then Map.delete gsz gs
                             else Map.insert gsz as gs
            delta = gsz  sz