RTree.hs

{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}

{- |
Module: Control.RTree
-}
module Control.RTree where

import Control.Applicative
import Control.Monad.Reader
import Control.Monad.Trans.Maybe
import qualified Data.List.NonEmpty as NE
import qualified Data.Map.Strict as Map
import Data.Maybe
import Data.Void
import GHC.IORef

class (Monad m) => MonadReduce l m | m -> l where
  split :: Maybe l -> m i -> m i -> m i

infixr 3 <|
infixl 3 |>

{-# INLINE (<|) #-}
(<|) :: (MonadReduce l r) => r i -> r i -> r i
r1 <| r2 = split Nothing r1 r2

{-# INLINE splitOn #-}
splitOn :: (MonadReduce l r) => l -> r i -> r i -> r i
splitOn l = split (Just l)

{-# INLINE (|>) #-}
(|>) :: (MonadReduce l r) => r i -> r i -> r i
r1 |> r2 = r2 <| r1

data RTree l i
  = Split (RTree l i) !(RTree l i)
  | SplitOn !l (RTree l i) !(RTree l i)
  | Done i
  deriving (Functor)

extract :: RTree l i -> i
extract = \case
  Split _ rhs -> extract rhs
  SplitOn _ _ rhs -> extract rhs
  Done v -> v

instance Applicative (RTree l) where
  pure = Done
  (<*>) = ap

instance Monad (RTree l) where
  ma >>= f = case ma of
    Done i -> f i
    Split lhs rhs ->
      Split (lhs >>= f) (rhs >>= f)
    SplitOn l lhs rhs ->
      SplitOn l (lhs >>= f) (rhs >>= f)

instance MonadReduce l (RTree l) where
  split = \case
    Just n -> SplitOn n
    Nothing -> Split

reduce
  :: forall m l i
   . (Alternative m)
  => (i -> m ())
  -> RTree l i
  -> m i
reduce p = checkgo
 where
  go = \case
    (Done i) -> pure i
    (Split lhs rhs) ->
      (checkgo lhs Control.Applicative.<|> go rhs)
    (SplitOn _ lhs rhs) ->
      (checkgo lhs Control.Applicative.<|> go rhs)
  checkgo rt = p (extract rt) *> go rt
{-# SPECIALIZE reduce :: (i -> MaybeT IO ()) -> RTree l i -> MaybeT IO i #-}

type Valuation l = Map.Map l Bool

extractL :: (Ord l) => Valuation l -> RTree l i -> i
extractL v = \case
  Split _ rhs -> extractL v rhs
  SplitOn l lhs rhs -> case Map.lookup l v of
    Just False -> extractL v lhs
    _ -> extractL v rhs
  Done i -> i

reduceL
  :: forall m l i
   . (Alternative m, Ord l)
  => (Valuation l -> i -> m ())
  -> Valuation l
  -> RTree l i
  -> m i
reduceL p = checkgo
 where
  checkgo v r = p v (extractL v r) *> go v r
  go v = \case
    Done i -> pure i
    SplitOn l lhs rhs -> case Map.lookup l v of
      Just True -> checkgo v rhs
      Just False -> checkgo v lhs
      Nothing -> checkgo (Map.insert l False v) lhs <|> go (Map.insert l True v) rhs
    Split lhs rhs -> (checkgo v lhs <|> go v rhs)
{-# INLINE reduceL #-}

data ReState l = ReState ![Bool] !(Valuation l)

newtype ReReduce l i = ReReduce {runReReduce :: IORef (ReState l) -> IO i}
  deriving (Functor, Applicative, Monad) via (ReaderT (IORef (ReState l)) IO)

instance (Ord l) => MonadReduce l (ReReduce l) where
  split ml r1 r2 = ReReduce \ref -> do
    test <- case ml of
      Nothing -> do
        atomicModifyIORef'
          ref
          ( \case
              ReState (a : as) v -> (ReState as v, a)
              ReState [] v -> (ReState [] v, False)
          )
      Just l -> do
        atomicModifyIORef'
          ref
          ( \case
              ReState as v@(Map.lookup l -> Just x) -> (ReState as v, not x)
              ReState (a : as) v -> (ReState as (Map.insert l (not a) v), a)
              ReState [] v -> (ReState [] (Map.insert l True v), False)
          )
    if test
      then runReReduce r1 ref
      else runReReduce r2 ref

reduceFast
  :: forall m l i
   . (MonadIO m, Ord l)
  => (Valuation l -> i -> MaybeT m ())
  -> Valuation l
  -> ReReduce l i
  -> MaybeT m i
reduceFast p v (ReReduce ext) = MaybeT $ go []
 where
  go pth = do
    ref <- liftIO $ newIORef (ReState pth v)
    i <- liftIO $ ext ref
    ReState r v' <- liftIO $ readIORef ref
    case r of
      [] -> do
        t <- isJust <$> runMaybeT (p v' i)
        if t
          then go (pth <> [True])
          else go (pth <> [False])
      _
        | null pth ->
            pure Nothing
        | otherwise ->
            pure (Just i)
{-# INLINE reduceFast #-}

-- Combinators

type MRTree l = MaybeT (RTree l)

instance (MonadReduce l m) => MonadReduce l (MaybeT m) where
  split m (MaybeT lhs) (MaybeT rhs) = MaybeT (split m lhs rhs)

-- | Given a list of item try to remove each of them the list.
collect :: (MonadReduce l m) => (a -> MaybeT m b) -> [a] -> m [b]
collect fn = fmap catMaybes . traverse (runMaybeT . fn)
{-# INLINE collect #-}

collectNonEmpty' :: (MonadReduce l m) => (a -> MaybeT m b) -> [a] -> MaybeT m [b]
collectNonEmpty' fn as =
  NE.toList <$> collectNonEmpty fn as
{-# INLINE collectNonEmpty' #-}

collectNonEmpty :: (MonadReduce l m) => (a -> MaybeT m b) -> [a] -> MaybeT m (NE.NonEmpty b)
collectNonEmpty fn as = do
  as' <- lift . fmap catMaybes . traverse (runMaybeT . fn) $ as
  MaybeT . pure $ NE.nonEmpty as'
{-# INLINE collectNonEmpty #-}

-- newtype LTree l i = LTree {runLTree :: Valuation l -> Maybe (RTree l i)}
--   deriving (Functor)
--
-- instance Applicative (LTree l) where
--   pure i = LTree{runLTree = \_ -> Just $ Done i}
--   (<*>) = ap
--
-- instance Monad (LTree l) where
--   LTree ma >>= f = LTree \l ->
--     case ma l of
--       Done i -> f i
--       Split l lhs rhs ->

-- extract' :: RTree l i -> i
-- extract' = \case
--   RTree' (Split _ _ v) -> extract' v
--   Done v -> v
--
-- instance Functor (RTree l) where
--   fmap f (Done i) = Done (f i)
--   fmap f (RTree' r) = RTree' (fmap (fmap f) r)
--
--
-- instance MonadFree (RTreeF l) (RTree' l) where
--   wrap = RTree'
--   {-# INLINE wrap #-}

-- | Reduce an input using a monad.

-- newtype I l i = I ([(l, Bool)] -> RTreeI l i)
--
-- data RTreeI l i
--   = RTreeI (RTreeF l (I l i))
--   | DoneI !i

-- -- This is not a great defintions, as the i does not depend on
-- -- the current i, but instead on the final I.
-- data RTreeIO j i = RTreeIO ((j -> IO Bool) -> IO i) j
--
-- extractIO :: RTreeIO j i -> j
-- extractIO (RTreeIO _ i) = i

-- instance Functor (RTreeIO j) where
--   fmap f (RTreeIO mf i) = RTreeIO (\h -> f <$> mf (h . f)) (f i)
--
-- instance Applicative (RTreeIO j) where
--   pure i = RTreeIO (\_ -> pure i) i
--   (<*>) = ap
--
-- -- RTreeIO f fi <*> RTreeIO a ai = RTreeIO (f <*> a) (fi ai)
--
-- instance Monad (RTreeIO j) where
--   RTreeIO (ma :: ((a -> IO Bool) -> IO a)) a >>= (f :: (a -> RTreeIO b)) =
--     RTreeIO undefined (extractIO $ f a)
--
-- instance MonadFree (RTreeF Void) (RTreeIO j) where
--   wrap (Split Nothing (RTreeIO lhs le) (RTreeIO rhs re)) =
--     RTreeIO
--       ( \p ->
--           p le >>= \case
--             True -> lhs p
--             False -> rhs p
--       )
--       re
--   wrap (Split (Just x) _ _) = absurd x

-- reduceIO
--   :: forall i
--    . (i -> IO Bool)
--   -> RTreeIO j i
--   -> IO (Maybe i)
-- reduceIO p (RTreeIO rt i) = runMaybeT do
--   let (mi, i') = foldR go $ fmap (\i -> (pure i, i)) t
--   p i' *> mi
--  where
--   go :: RTreeF l (IO i, i) -> (IO i, i)
--   go (Split _ (lhs, le) (rhs, re)) =
--     ((p le *> lhs) <|> rhs, re)

-- | Split the world on a fact. False it does not happen, and True it does happen.
given :: RTree Void Bool
given = pure False <| pure True

{- | A reducer should extract itself
@
 extract . red = id
@
-}
lawReduceId :: (Eq i) => (i -> RTree l i) -> i -> Bool
lawReduceId red i = extract (red i) == i

-- | Reducing a list one element at a time.
rList :: [a] -> RTree l [a]
rList = \case
  [] -> pure []
  a : as -> rList as <| (a :) <$> rList as

{- | Binary reduction on the list assumming suffixes all contain eachother:
@[] < [c] < [b, c] < [a,b,c]@
-}
rSuffixList :: [a] -> RTree l [a]
rSuffixList as = do
  res <- exponentialSearch (NE.tails as)
  case res of
    [] -> pure []
    a : as' -> (a :) <$> rSuffixList as'

{- | Given a progression of inputs that are progressively larger, pick the smallest using
binary search.
-}
binarySearch :: NE.NonEmpty i -> RTree l i
binarySearch = \case
  a NE.:| [] -> pure a
  d -> binarySearch l <| binarySearch f
   where
    (NE.fromList -> f, NE.fromList -> l) = NE.splitAt (NE.length d `div` 2) d

{- | Given a progression of inputs that are progressively larger, pick the smallest using
binary search.
-}
exponentialSearch :: NE.NonEmpty i -> RTree l i
exponentialSearch = go 1
 where
  go n = \case
    d
      | n >= NE.length d -> binarySearch d
      | otherwise -> go (n * 2) l <| binarySearch f
     where
      (NE.fromList -> f, NE.fromList -> l) = NE.splitAt n d

nonEmptyOr :: String -> [a] -> NE.NonEmpty a
nonEmptyOr msg ls = case NE.nonEmpty ls of
  Just a -> a
  Nothing -> error msg

-- | Given a list of orderd options,  the
linearSearch :: NE.NonEmpty i -> RTree l i
linearSearch = foldr1 (<|) . fmap pure

-- | Given a list of orderd options,  the
linearSearch' :: [i] -> RTree l (Maybe i)
linearSearch' is = linearSearch (NE.fromList $ fmap Just is ++ [Nothing])

-- | Given
ddmin :: [i] -> RTree l [i]
ddmin = \case
  [] -> pure []
  [a] -> pure [a]
  as -> go 2 as
 where
  go n lst
    | n' <= 0 = pure lst
    | otherwise = do
        r <- linearSearch' (partitions n' lst ++ composites n' lst)
        case r of
          Nothing -> go (n * 2) lst <| pure lst -- (for efficiency :D)
          Just lst' -> ddmin lst'
   where
    n' = length lst `div` n
  partitions n lst =
    case lst of
      [] -> []
      _ -> let (h, r) = splitAt n lst in h : partitions n r
  composites n lst =
    case lst of
      [] -> []
      _ -> let (h, r) = splitAt n lst in r : fmap (h ++) (composites n r)