content: {-# LANGUAGE DeriveGeneric #-}
  {-# LANGUAGE DeriveFunctor #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE TemplateHaskell #-}
  {-# LANGUAGE UndecidableInstances #-}
  {-# LANGUAGE StandaloneDeriving #-}
  {-# LANGUAGE RecordWildCards #-}
  
  module BTree.Types where
  
  import Control.Applicative
  import Data.Maybe (fromMaybe)
  import GHC.Generics
  import Control.Monad (when)
  import Data.Int
  import Prelude
  
  import Data.Binary
  import Data.Binary.Get
  import Data.Binary.Put
  import Control.Lens
  import qualified Data.Vector as V
  import Data.Vector.Binary
  import qualified Data.ByteString as BS
  
  -- | An offset within the stream
  type Offset = Int64
  
  -- | The number of entries in a B-tree
  type Size = Word64
  
  -- | The maximum number of children of a B-tree inner node
  type Order = Word64
  
  -- | @'OnDisk' a@ is a reference to an object of type @a@ on disk.
  -- The offset does not include the header; e.g. the first object after
  -- the header is located at offset 0.
  newtype OnDisk a = OnDisk Offset
                   deriving (Show, Eq, Ord)
  
  instance Binary (OnDisk a) where
      get = OnDisk <$> get
      {-# INLINE get #-}
      put (OnDisk off) = put off
      {-# INLINE put #-}
  
  -- | A tree leaf (e.g. key/value pair)
  data BLeaf k e = BLeaf !k !e
                 deriving (Generic, Functor)
  
  deriving instance (Show k, Show e) => Show (BLeaf k e)
  
  -- | This only compares on the keys
  instance (Eq k) => Eq (BLeaf k e) where
      BLeaf a _ == BLeaf b _ = a == b
  
  -- | This only compares on the keys
  instance Ord k => Ord (BLeaf k e) where
      compare (BLeaf a _) (BLeaf b _) = compare a b
      {-# INLINE compare #-}
  
  instance (Binary k, Binary e) => Binary (BLeaf k e) where
      get = BLeaf <$> get <*> get
      {-# INLINE get #-}
      put (BLeaf k e) = put k >> put e
      {-# INLINE put #-}
  
  -- | @'BTree' k f e@ is a B* tree of key type @k@ with elements of type @e@.
  -- Subtree references are contained within a type @f@.
  --
  -- The 'Node' constructor contains a left child, and a list of key/child pairs
  -- where each child's keys are greater than or equal to the given key.
  data BTree k f e = Node (f (BTree k f e)) (V.Vector (k, f (BTree k f e)))
                   | Leaf !(BLeaf k e)
                   deriving (Generic)
  
  deriving instance (Show e, Show k, Show (f (BTree k f e))) => Show (BTree k f e)
  deriving instance (Eq e, Eq k, Eq (f (BTree k f e))) => Eq (BTree k f e)
  
  instance (Binary k, Binary (f (BTree k f e)), Binary e)
    => Binary (BTree k f e) where
      get = do typ <- getWord8
               case typ of
                 0 -> Node <$> get <*> getChildren
                 1 -> bleaf <$> get <*> get
                 _ -> fail "BTree.Types/get: Unknown node type"
        where bleaf k v = Leaf (BLeaf k v)
              getChildren =
                  genericGetVectorWith (fromIntegral <$> getWord32be) ((,) <$> get <*> get)
      {-# INLINE get #-}
  
      -- some versions of binary don't inline the Binary (,) instance, pitiful
      -- performance ensues
      put (Node e0 es) = do
          putWord8 0
          put e0
          genericPutVectorWith (putWord32be . fromIntegral) (\(a,b) -> put a >> put b) es
      put (Leaf (BLeaf k0 e)) = putWord8 1 >> put k0 >> put e
      {-# INLINE put #-}
  
  magic :: Word64
  magic = 0xdeadbeefbbbbcccc
  
  -- | B-tree file header
  data BTreeHeader k e = BTreeHeader { _btMagic   :: !Word64
                                     , _btVersion :: !Word64
                                     , _btOrder   :: !Order
                                     , _btSize    :: !Size
                                     , _btRoot    :: !(Maybe (OnDisk (BTree k OnDisk e)))
                                       -- ^ 'Nothing' represents an empty tree
                                     }
                   deriving (Show, Eq, Generic)
  makeLenses ''BTreeHeader
  
  -- | It is critical that this encoding is of fixed size
  instance Binary (BTreeHeader k e) where
      get = do
          _btMagic <- get
          _btVersion <- get
          _btOrder <- get
          _btSize <- get
          root <- get
          let _btRoot = if root == OnDisk 0 then Nothing else Just root
          return BTreeHeader {..}
      put (BTreeHeader {..}) = do
          put _btMagic
          put _btVersion
          put _btOrder
          put _btSize
          put $ fromMaybe (OnDisk 0) _btRoot
  
  validateHeader :: BTreeHeader k e -> Either String ()
  validateHeader hdr = do
      when (hdr^.btMagic /= magic) $ Left "Invalid magic number"
      when (hdr^.btVersion > 1) $ Left "Invalid version"
  
  -- | A read-only B-tree for lookups
  data LookupTree k e = LookupTree { _ltData    :: !BS.ByteString
                                   , _ltHeader  :: !(BTreeHeader k e)
                                   }
  makeLenses ''LookupTree
  