{-# LANGUAGE FlexibleInstances #-}

module HaskellWorks.Data.BalancedParens.BalancedParens
  ( BalancedParens(..)
  , depth
  , subtreeSize
  ) where

import Control.Monad
import Data.Word
import HaskellWorks.Data.BalancedParens.CloseAt
import HaskellWorks.Data.BalancedParens.Enclose
import HaskellWorks.Data.BalancedParens.FindClose
import HaskellWorks.Data.BalancedParens.FindOpen
import HaskellWorks.Data.BalancedParens.OpenAt
import HaskellWorks.Data.Naive
import HaskellWorks.Data.Positioning
import HaskellWorks.Data.RankSelect.Base.Rank0
import HaskellWorks.Data.RankSelect.Base.Rank1

import qualified Data.Vector.Storable as DVS

class (OpenAt v, CloseAt v, FindOpen v, FindClose v, Enclose v) => BalancedParens v where
  -- TODO Second argument should be Int
  firstChild :: v -> Count -> Maybe Count
  firstChild  v
v Count
p = if v -> Count -> Bool
forall v. OpenAt v => v -> Count -> Bool
openAt v
v Count
p Bool -> Bool -> Bool
&& v -> Count -> Bool
forall v. OpenAt v => v -> Count -> Bool
openAt v
v (Count
p Count -> Count -> Count
forall a. Num a => a -> a -> a
+ Count
1)   then Count -> Maybe Count
forall a. a -> Maybe a
Just (Count
p Count -> Count -> Count
forall a. Num a => a -> a -> a
+ Count
1) else Maybe Count
forall a. Maybe a
Nothing
  {-# INLINE firstChild #-}

  nextSibling :: v -> Count -> Maybe Count
  nextSibling v
v Count
p = if v -> Count -> Bool
forall v. CloseAt v => v -> Count -> Bool
closeAt v
v Count
p
    then Maybe Count
forall a. Maybe a
Nothing
    else v -> Count -> Bool
forall v. OpenAt v => v -> Count -> Bool
openAt v
v (Count -> Bool) -> Maybe Count -> Maybe Count
forall (m :: * -> *) a. MonadPlus m => (a -> Bool) -> m a -> m a
`mfilter` (v -> Count -> Maybe Count
forall v. FindClose v => v -> Count -> Maybe Count
findClose v
v Count
p Maybe Count -> (Count -> Maybe Count) -> Maybe Count
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (\Count
q ->
      if Count
p Count -> Count -> Bool
forall a. Eq a => a -> a -> Bool
/= Count
q
        then Count -> Maybe Count
forall (m :: * -> *) a. Monad m => a -> m a
return (Count
q Count -> Count -> Count
forall a. Num a => a -> a -> a
+ Count
1)
        else Maybe Count
forall a. Maybe a
Nothing))
  {-# INLINE nextSibling #-}

  parent :: v -> Count -> Maybe Count
  parent v
v Count
p = v -> Count -> Maybe Count
forall v. Enclose v => v -> Count -> Maybe Count
enclose v
v Count
p Maybe Count -> (Count -> Maybe Count) -> Maybe Count
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (\Count
r -> if Count
r Count -> Count -> Bool
forall a. Ord a => a -> a -> Bool
>= Count
1 then Count -> Maybe Count
forall (m :: * -> *) a. Monad m => a -> m a
return Count
r      else Maybe Count
forall a. Maybe a
Nothing)
  {-# INLINE parent #-}

depth :: (BalancedParens v, Rank0 v, Rank1 v) => v -> Count -> Maybe Count
depth :: v -> Count -> Maybe Count
depth v
v Count
p = (\Count
q -> v -> Count -> Count
forall v. Rank1 v => v -> Count -> Count
rank1 v
v Count
q Count -> Count -> Count
forall a. Num a => a -> a -> a
- v -> Count -> Count
forall v. Rank0 v => v -> Count -> Count
rank0 v
v Count
q) (Count -> Count) -> Maybe Count -> Maybe Count
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> v -> Count -> Maybe Count
forall v. FindOpen v => v -> Count -> Maybe Count
findOpen v
v Count
p
{-# INLINE depth #-}

subtreeSize :: BalancedParens v => v -> Count -> Maybe Count
subtreeSize :: v -> Count -> Maybe Count
subtreeSize v
v Count
p = (\Count
q -> (Count
q Count -> Count -> Count
forall a. Num a => a -> a -> a
- Count
p Count -> Count -> Count
forall a. Num a => a -> a -> a
+ Count
1) Count -> Count -> Count
forall a. Integral a => a -> a -> a
`quot` Count
2) (Count -> Count) -> Maybe Count -> Maybe Count
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> v -> Count -> Maybe Count
forall v. FindClose v => v -> Count -> Maybe Count
findClose v
v Count
p
{-# INLINE subtreeSize #-}

instance BalancedParens [Bool]

instance BalancedParens (DVS.Vector Word8)

instance BalancedParens (DVS.Vector Word16)

instance BalancedParens (DVS.Vector Word32)

instance BalancedParens (DVS.Vector Word64)

instance BalancedParens Word8

instance BalancedParens Word16

instance BalancedParens Word32

instance BalancedParens Word64

instance BalancedParens (Naive Word64)