content: module BTree.Lookup ( LookupTree
                      , open
                      , fromByteString
                      , lookup
                      , size
                      ) where
  
  import Prelude hiding (lookup)
  import Control.Error
  import Control.Lens hiding (children)
  import qualified Data.ByteString as BS
  import qualified Data.Vector as V
  import qualified Data.ByteString.Lazy as LBS
  import Data.Binary
  import System.IO.MMap
  import BTree.Types
  
  fetch :: (Binary a) => LookupTree k e -> OnDisk a -> a
  fetch lt (OnDisk offset) =
      decode $ LBS.fromStrict $ BS.drop (fromIntegral offset) (lt^.ltData)
  
  -- | Read a B-tree from a 'ByteString' produced by 'BTree.Builder'
  fromByteString :: LBS.ByteString -> Either String (LookupTree k e)
  fromByteString bs = do
      (rest, _, hdr) <- fmapL (\(_,_,e)->e) $ decodeOrFail bs
      validateHeader hdr
      return $ LookupTree (LBS.toStrict rest) hdr
  
  -- | Open a B-tree file.
  open :: FilePath -> IO (Either String (LookupTree k e))
  open fname = runExceptT $ do
      d <- fmapLT show $ tryIO $ mmapFileByteString fname Nothing
      ExceptT $ return $ fromByteString (LBS.fromStrict d)
  
  -- | Lookup a key in a B-tree.
  lookup :: (Binary k, Binary e, Ord k)
         => LookupTree k e -> k -> Maybe e
  lookup lt k =
      case lt ^. ltHeader . btRoot of
        Just root -> go $ fetch lt root
        Nothing   -> Nothing
    where
      go (Leaf (BLeaf k' e))
        | k' == k     = Just e
        | otherwise   = Nothing
      go (Node c0 children)
        | V.null children = go $ fetch lt c0 -- is this case necessary?
        | let (k0,_) = V.head children
        , k < k0      = go $ fetch lt c0
        | otherwise   =
            case V.takeWhile (\(k',_)->k' <= k) children of
              rest
                | V.null rest -> Nothing
                | otherwise   -> go $ fetch lt $ snd $ V.last rest
  
  -- | How many keys are in a 'LookupTree'.
  size :: LookupTree k e -> Size
  size = _btSize . _ltHeader
  