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 <|> go rhs)
    (SplitOn _ lhs rhs) -> (checkgo lhs <|> 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)

{- | A faster version of the RTree which simply reruns the reducer program instead
of building a tree.
-}
newtype IORTree l i = IORTree {runIORTree :: IORef (ReState l) -> IO i}
  deriving (Functor, Applicative, Monad) via (ReaderT (IORef (ReState l)) IO)

instance (Ord l) => MonadReduce l (IORTree l) where
  split ml r1 r2 = IORTree \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 runIORTree r1 ref
      else runIORTree r2 ref

reduceIO
  :: forall m l i
   . (MonadIO m, Ord l)
  => (Valuation l -> i -> MaybeT m ())
  -> Valuation l
  -> IORTree l i
  -> MaybeT m i
reduceIO p v (IORTree 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 reduceIO #-}

-- 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 #-}

-- | 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)