content: {-
    Copyright (c) Meta Platforms, Inc. and affiliates.
    All rights reserved.
  
    This source code is licensed under the BSD-style license found in the
    LICENSE file in the root directory of this source tree.
  -}
  
  module Glean.Util.Bisect (
    Bisect(..), BisectM, execBisectM, execBisectRefM, bisect
  ) where
  
  import Control.Applicative
  import Control.Monad
  import Control.Monad.IO.Class
  import Control.Monad.Trans.Except (ExceptT, runExceptT, throwE)
  import Control.Monad.Trans.Reader
  import Control.Monad.Trans.Class (lift)
  import Data.IORef
  
  data Bisect a b err = Bisect
    { bisectProcess :: [a] -> IO (Either err b)
    , bisectFailed :: err -> a -> IO ()
      -- ^ called if the bisect successfully identified a failing item,
      -- regardless of whether we continued or aborted afterwards.
    , bisectSplitting :: Int -> Int -> [a] -> IO ()
    }
  
  newtype BisectM err a = BisectM (ExceptT err (ReaderT (IORef Int) IO) a)
    deriving(Functor, Applicative, Alternative, Monad, MonadPlus, MonadIO)
  
  -- | Run a bisect operation
  execBisectM
    :: Int
       -- ^ Number of failures that we will bisect down the failing item
       -- and continue.
    -> BisectM err a
       -- ^ e.g. a call to 'bisect'
    -> IO (Either err a)
  execBisectM allowed (BisectM action) = do
    ref <- newIORef allowed
    runReaderT (runExceptT action) ref
  
  -- | Run a bisect operation
  execBisectRefM
    :: IORef Int
       -- ^ Number of failures that we will bisect down the failing item
       -- and continue.
    -> BisectM err a
       -- ^ e.g. a call to 'bisect'
    -> IO (Either err a)
  execBisectRefM ref (BisectM action) = runReaderT (runExceptT action) ref
  
  bisect
    :: Bisect a b err
    -> [a]
    -> BisectM (err,[a]) [(b, [a])]
       -- ^ On abort, returns the last error and the list of items
  bisect !bis xs = BisectM $ attempt 0 (length xs) xs
    where
      attempt k n xs = do
        r <- liftIO $ bisectProcess bis xs
        case r of
          Right x -> return [(x, xs)]
          Left err -> do
            r <- lift ask
            case xs of
              [x] -> do
                liftIO $ bisectFailed bis err x
                e <- liftIO $ atomicModifyIORef' r (\x -> (max 0 (x-1), x))
                when (e == 0) $ throwE (err,xs)
                return []
              _otherwise -> do
                e <- liftIO $ readIORef r
                if e > 0
                  then split k n xs
                  else throwE (err,xs)
  
      split _ _ [] = return []
      split k n xs = do
        liftIO $ bisectSplitting bis k n xs
        (++) <$> attempt k m ys <*> attempt (k+m) (n-m) zs
        where
          m = n `div` 2
          (ys,zs) = splitAt m xs
  