{-# LANGUAGE BangPatterns #-}

{-| Spine-strict radix tree location manipulation.

    Allows lookup and successive insertion without retaining the entirety
    of the key in memory.
 -}

module Data.Radix1Tree.Word8.Strict.Zipper
  ( Context1
  , descend
  , focus
  ) where

import           Data.ByteArray.NonEmpty
import           Data.RadixNTree.Word8.Key
import           Data.RadixNTree.Word8.Strict
import           Radix.Word8.Foundation

import           Data.Primitive.ByteArray



data Past a = Leftward
                !(Past a)
                {-# UNPACK #-} !Prefix
                !(Radix1Tree a)

            | Rightward
                !(Past a)
                {-# UNPACK #-} !Prefix
                !(Radix1Tree a)

            | Downward
                !(Past a)
                {-# UNPACK #-} !ByteArray
                {-# UNPACK #-} !(Maybe a)

            | Top


-- | A location inside the radix tree.
data Context1 a = -- | Corresponds to a 'Tip'.
                  Context1
                    !(Past a)
                    {-# UNPACK #-} !Int       -- ^ Next index in the byte array.
                    {-# UNPACK #-} !ByteArray
                    {-# UNPACK #-} !(Maybe a)
                    !(Radix1Tree a)



{-# INLINE descend #-}
-- | \(\mathcal{O}(\min(x,k))\).
--   Move down the tree by the extent of the given key.
--   Returns 'Nothing' if the resulting position is outside of the tree.
--
--   @since 1.1
descend :: Feed1 -> Either (Radix1Tree a) (Context1 a) -> Maybe (Context1 a)
descend :: forall a.
Feed1 -> Either (Radix1Tree a) (Context1 a) -> Maybe (Context1 a)
descend (Feed1 Word8
w0 forall a. (forall x. (x -> Step Word8 x) -> x -> a) -> a
feed) =
  (forall x.
 (x -> Step Word8 x)
 -> x -> Either (Radix1Tree a) (Context1 a) -> Maybe (Context1 a))
-> Either (Radix1Tree a) (Context1 a) -> Maybe (Context1 a)
forall a. (forall x. (x -> Step Word8 x) -> x -> a) -> a
feed ((forall x.
  (x -> Step Word8 x)
  -> x -> Either (Radix1Tree a) (Context1 a) -> Maybe (Context1 a))
 -> Either (Radix1Tree a) (Context1 a) -> Maybe (Context1 a))
-> (forall x.
    (x -> Step Word8 x)
    -> x -> Either (Radix1Tree a) (Context1 a) -> Maybe (Context1 a))
-> Either (Radix1Tree a) (Context1 a)
-> Maybe (Context1 a)
forall a b. (a -> b) -> a -> b
$ \x -> Step Word8 x
step ->

    let go :: Past a -> Word8 -> x -> Radix1Tree a -> Maybe (Context1 a)
go !Past a
past !Word8
w !x
s Radix1Tree a
t =
          case Radix1Tree a
t of
            Bin Word8
p Radix1Tree a
l Radix1Tree a
r ->
              if Word8
w Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
< Word8
p
                then Past a -> Word8 -> x -> Radix1Tree a -> Maybe (Context1 a)
go (Past a -> Word8 -> Radix1Tree a -> Past a
forall a. Past a -> Word8 -> Radix1Tree a -> Past a
Leftward  Past a
past Word8
p Radix1Tree a
r) Word8
w x
s Radix1Tree a
l
                else Past a -> Word8 -> x -> Radix1Tree a -> Maybe (Context1 a)
go (Past a -> Word8 -> Radix1Tree a -> Past a
forall a. Past a -> Word8 -> Radix1Tree a -> Past a
Rightward Past a
past Word8
p Radix1Tree a
l) Word8
w x
s Radix1Tree a
r

            Tip ByteArray
arr Maybe a
mx Radix1Tree a
dx -> Past a
-> ByteArray
-> Maybe a
-> Radix1Tree a
-> Word8
-> x
-> Int
-> Maybe (Context1 a)
goarr Past a
past ByteArray
arr Maybe a
mx Radix1Tree a
dx Word8
w x
s Int
0

            Radix1Tree a
Nil -> Maybe (Context1 a)
forall a. Maybe a
Nothing

        goarr :: Past a
-> ByteArray
-> Maybe a
-> Radix1Tree a
-> Word8
-> x
-> Int
-> Maybe (Context1 a)
goarr !Past a
past !ByteArray
arr !Maybe a
mx Radix1Tree a
dx = Word8 -> x -> Int -> Maybe (Context1 a)
goarr_
          where
            goarr_ :: Word8 -> x -> Int -> Maybe (Context1 a)
goarr_ Word8
v !x
z Int
n
              | Word8
v Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== ByteArray -> Int -> Word8
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
arr Int
n =
                  let n' :: Int
n' = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
                  in case x -> Step Word8 x
step x
z of
                       More Word8
u x
z' -> if Int
n' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= ByteArray -> Int
sizeofByteArray ByteArray
arr
                                      then Past a -> Word8 -> x -> Radix1Tree a -> Maybe (Context1 a)
go (Past a -> ByteArray -> Maybe a -> Past a
forall a. Past a -> ByteArray -> Maybe a -> Past a
Downward Past a
past ByteArray
arr Maybe a
mx) Word8
u x
z' Radix1Tree a
dx
                                      else Word8 -> x -> Int -> Maybe (Context1 a)
goarr_ Word8
u x
z' Int
n'

                       Step Word8 x
Done      -> Context1 a -> Maybe (Context1 a)
forall a. a -> Maybe a
Just (Context1 a -> Maybe (Context1 a))
-> Context1 a -> Maybe (Context1 a)
forall a b. (a -> b) -> a -> b
$! Past a -> Int -> ByteArray -> Maybe a -> Radix1Tree a -> Context1 a
forall a.
Past a -> Int -> ByteArray -> Maybe a -> Radix1Tree a -> Context1 a
Context1 Past a
past Int
n' ByteArray
arr Maybe a
mx Radix1Tree a
dx

              | Bool
otherwise = Maybe (Context1 a)
forall a. Maybe a
Nothing

    in \x
s0 Either (Radix1Tree a) (Context1 a)
ei ->
         case Either (Radix1Tree a) (Context1 a)
ei of
           Left Radix1Tree a
r                           -> Past a -> Word8 -> x -> Radix1Tree a -> Maybe (Context1 a)
forall {a}.
Past a -> Word8 -> x -> Radix1Tree a -> Maybe (Context1 a)
go Past a
forall a. Past a
Top Word8
w0 x
s0 Radix1Tree a
r
           Right (Context1 Past a
past Int
n ByteArray
arr Maybe a
mx Radix1Tree a
dx) ->
             if Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== ByteArray -> Int
sizeofByteArray ByteArray
arr
               then Past a -> Word8 -> x -> Radix1Tree a -> Maybe (Context1 a)
forall {a}.
Past a -> Word8 -> x -> Radix1Tree a -> Maybe (Context1 a)
go (Past a -> ByteArray -> Maybe a -> Past a
forall a. Past a -> ByteArray -> Maybe a -> Past a
Downward Past a
past ByteArray
arr Maybe a
mx) Word8
w0 x
s0 Radix1Tree a
dx
               else Past a
-> ByteArray
-> Maybe a
-> Radix1Tree a
-> Word8
-> x
-> Int
-> Maybe (Context1 a)
forall {a}.
Past a
-> ByteArray
-> Maybe a
-> Radix1Tree a
-> Word8
-> x
-> Int
-> Maybe (Context1 a)
goarr Past a
past ByteArray
arr Maybe a
mx Radix1Tree a
dx Word8
w0 x
s0 Int
n



-- | \(\mathcal{O}(1)\).
--   Retrieve the value at the current position, if any exists,
--   together with the insertion function for the current position.
--
--   @since 1.1
focus :: Context1 a -> Maybe (a, a -> Radix1Tree a)
focus :: forall a. Context1 a -> Maybe (a, a -> Radix1Tree a)
focus (Context1 Past a
past Int
n ByteArray
arr Maybe a
mx Radix1Tree a
dx)
  | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== ByteArray -> Int
sizeofByteArray ByteArray
arr, Just a
x <- Maybe a
mx =
      (a, a -> Radix1Tree a) -> Maybe (a, a -> Radix1Tree a)
forall a. a -> Maybe a
Just ((a, a -> Radix1Tree a) -> Maybe (a, a -> Radix1Tree a))
-> (a, a -> Radix1Tree a) -> Maybe (a, a -> Radix1Tree a)
forall a b. (a -> b) -> a -> b
$! (a
x, \a
y -> Radix1Tree a -> Past a -> Radix1Tree a
forall a. Radix1Tree a -> Past a -> Radix1Tree a
rebuild (ByteArray -> Maybe a -> Radix1Tree a -> Radix1Tree a
forall a. ByteArray -> Maybe a -> Radix1Tree a -> Radix1Tree a
Tip ByteArray
arr (a -> Maybe a
forall a. a -> Maybe a
Just a
y) Radix1Tree a
dx) Past a
past)

  | Bool
otherwise = Maybe (a, a -> Radix1Tree a)
forall a. Maybe a
Nothing



rebuild :: Radix1Tree a -> Past a -> Radix1Tree a
rebuild :: forall a. Radix1Tree a -> Past a -> Radix1Tree a
rebuild !Radix1Tree a
x Past a
past =
  case Past a
past of
    Leftward  Past a
past' Word8
p Radix1Tree a
r   -> Radix1Tree a -> Past a -> Radix1Tree a
forall a. Radix1Tree a -> Past a -> Radix1Tree a
rebuild (Word8 -> Radix1Tree a -> Radix1Tree a -> Radix1Tree a
forall a. Word8 -> Radix1Tree a -> Radix1Tree a -> Radix1Tree a
Bin Word8
p Radix1Tree a
x Radix1Tree a
r) Past a
past'
    Rightward Past a
past' Word8
p Radix1Tree a
l   -> Radix1Tree a -> Past a -> Radix1Tree a
forall a. Radix1Tree a -> Past a -> Radix1Tree a
rebuild (Word8 -> Radix1Tree a -> Radix1Tree a -> Radix1Tree a
forall a. Word8 -> Radix1Tree a -> Radix1Tree a -> Radix1Tree a
Bin Word8
p Radix1Tree a
l Radix1Tree a
x) Past a
past'
    Downward Past a
past' ByteArray
brr Maybe a
my -> Radix1Tree a -> Past a -> Radix1Tree a
forall a. Radix1Tree a -> Past a -> Radix1Tree a
rebuild (ByteArray -> Maybe a -> Radix1Tree a -> Radix1Tree a
forall a. ByteArray -> Maybe a -> Radix1Tree a -> Radix1Tree a
Tip ByteArray
brr Maybe a
my Radix1Tree a
x) Past a
past'
    Past a
Top                   -> Radix1Tree a
x