content: {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  
  -- |
  -- Module      :   Grisette.Lib.Data.List
  -- Copyright   :   (c) Sirui Lu 2021-2023
  -- License     :   BSD-3-Clause (see the LICENSE file)
  --
  -- Maintainer  :   siruilu@cs.washington.edu
  -- Stability   :   Experimental
  -- Portability :   GHC only
  module Grisette.Lib.Data.List
    ( -- * Special folds
      symAnd,
      symOr,
      symAny,
      symAll,
      mrgMaximum,
      symMaximum,
      mrgMinimum,
      symMinimum,
  
      -- * Sublists
  
      -- ** Extracting sublists
      mrgTake,
      mrgDrop,
      mrgSplitAt,
      mrgTakeWhile,
      mrgDropWhile,
      mrgDropWhileEnd,
      mrgSpan,
      mrgBreak,
      mrgStripPrefix,
      mrgGroup,
  
      -- ** Predicates
      symIsPrefixOf,
      symIsSuffixOf,
      symIsInfixOf,
      symIsSubsequenceOf,
  
      -- * Searching lists
  
      -- ** Searching by equality
      symElem,
      symNotElem,
      mrgLookup,
  
      -- ** Searching with a predicate
      mrgFind,
      mrgFilter,
      mrgPartition,
  
      -- * Indexing lists
      (.!?),
      mrgElemIndex,
      mrgElemIndices,
      mrgFindIndex,
      mrgFindIndices,
  
      -- * Special lists
  
      -- ** "Set" operations
      mrgNub,
      mrgDelete,
      (.\\),
      mrgUnion,
      mrgIntersect,
  
      -- ** Ordered lists (sorting not supported yet)
      mrgInsert,
  
      -- * Generalized functions
  
      -- ** The "By" operations
  
      -- *** User-supplied equality (replacing an 'SymEq' context)
      mrgNubBy,
      mrgDeleteBy,
      mrgDeleteFirstsBy,
      mrgUnionBy,
      mrgIntersectBy,
      mrgGroupBy,
  
      -- *** User-supplied comparison (replacing an 'SymOrd' context)
      mrgInsertBy,
      mrgMaximumBy,
      symMaximumBy,
      mrgMinimumBy,
      symMinimumBy,
    )
  where
  
  import Data.Bifunctor (Bifunctor (first, second))
  import Data.List (tails)
  import Data.Maybe (listToMaybe)
  import Grisette.Internal.Core.Control.Monad.Class.Union (MonadUnion)
  import Grisette.Internal.Core.Control.Monad.Union (Union)
  import Grisette.Internal.Core.Data.Class.ITEOp (ITEOp (symIte))
  import Grisette.Internal.Core.Data.Class.LogicalOp (LogicalOp (symNot, (.&&), (.||)))
  import Grisette.Internal.Core.Data.Class.Mergeable (Mergeable)
  import Grisette.Internal.Core.Data.Class.SimpleMergeable (SymBranching, mrgIf)
  import Grisette.Internal.Core.Data.Class.Solvable (Solvable (con))
  import Grisette.Internal.Core.Data.Class.SymEq (SymEq ((./=), (.==)))
  import Grisette.Internal.Core.Data.Class.SymOrd (SymOrd ((.<=), (.>=)))
  import Grisette.Internal.Core.Data.Class.UnionView (liftUnion)
  import Grisette.Internal.SymPrim.SymBool (SymBool)
  import Grisette.Lib.Control.Applicative (mrgPure)
  import Grisette.Lib.Control.Monad (mrgReturn)
  import Grisette.Lib.Data.Foldable
    ( mrgFind,
      mrgFoldlM,
      mrgMaximum,
      mrgMaximumBy,
      mrgMinimum,
      mrgMinimumBy,
      symAll,
      symAnd,
      symAny,
      symElem,
      symMaximum,
      symMaximumBy,
      symMinimum,
      symMinimumBy,
      symNotElem,
      symOr,
    )
  import Grisette.Lib.Data.Functor (mrgFmap)
  
  symListOpOnSymInt ::
    (Applicative u, SymBranching u, Mergeable b, Num int, SymOrd int) =>
    Bool ->
    (Int -> [a] -> b) ->
    int ->
    [a] ->
    u b
  symListOpOnSymInt reversed f x vs = do
    let zipped =
          (\n -> (fromIntegral n, mrgPure $ f n vs))
            <$> (if reversed then reverse else id) [1 .. length vs - 1]
    let outerMostX = if reversed then length vs else 0
    let innerMostX = if reversed then 0 else length vs
    let guardCond =
          if reversed then (x .>= fromIntegral (length vs)) else (x .<= 0)
    mrgIf guardCond (mrgPure $ f outerMostX vs) $
      foldr
        (\(n, l) acc -> mrgIf (x .== n) l acc)
        (mrgPure $ f innerMostX vs)
        zipped
  
  -- | Symbolic version of 'Data.List.take', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- Can generate O(n) cases and O(n) sized branch constraints.
  mrgTake ::
    (Applicative u, SymBranching u, Mergeable a, Num int, SymOrd int) =>
    int ->
    [a] ->
    u [a]
  mrgTake = symListOpOnSymInt False take
  
  -- | Symbolic version of 'Data.List.drop', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- Can generate O(n) cases and O(n) sized branch constraints.
  mrgDrop ::
    (Applicative u, SymBranching u, Mergeable a, Num int, SymOrd int) =>
    int ->
    [a] ->
    u [a]
  mrgDrop = symListOpOnSymInt True drop
  
  -- | Symbolic version of 'Data.List.splitAt', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- Can generate O(n) cases and O(n) sized branch constraints.
  mrgSplitAt ::
    forall a int u.
    (MonadUnion u, Mergeable a, Num int, SymOrd int) =>
    int ->
    [a] ->
    u ([a], [a])
  mrgSplitAt = symListOpOnSymInt False splitAt
  
  -- | Symbolic version of 'Data.List.takeWhile', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- Can generate O(n) cases and O(n) sized branch constraints.
  mrgTakeWhile ::
    (Applicative u, SymBranching u, Mergeable a) =>
    (a -> SymBool) ->
    [a] ->
    u [a]
  mrgTakeWhile _ [] = mrgPure []
  mrgTakeWhile p (x : xs) =
    mrgIf (p x) (mrgFmap (x :) $ mrgTakeWhile p xs) (mrgPure [])
  
  -- | Symbolic version of 'Data.List.dropWhile', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- Can generate O(n) cases and O(n) sized branch constraints.
  mrgDropWhile ::
    (Applicative u, SymBranching u, Mergeable a) =>
    (a -> SymBool) ->
    [a] ->
    u [a]
  mrgDropWhile _ [] = mrgPure []
  mrgDropWhile p r = do
    let allConds = reverse $ scanl1 (.&&) $ p <$> r
    foldr (\(cond, l) acc -> mrgIf cond (pure l) acc) (pure r) $
      zip allConds $
        reverse $
          tails r
  
  -- | Symbolic version of 'Data.List.dropWhileEnd', the result would be merged
  -- and propagate the mergeable knowledge.
  --
  -- Can generate O(n) cases and O(n) sized branch constraints.
  mrgDropWhileEnd ::
    (MonadUnion u, Mergeable a) =>
    (a -> SymBool) ->
    [a] ->
    u [a]
  mrgDropWhileEnd p =
    foldr
      ( \x xs -> do
          xsv <- xs
          mrgIf (p x .&& con (null xsv)) (mrgPure []) (mrgPure $ x : xsv)
      )
      (mrgPure [])
  
  -- | Symbolic version of 'Data.List.span', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- Can generate O(n) cases and O(n) sized branch constraints.
  mrgSpan ::
    (Applicative u, SymBranching u, Mergeable a) =>
    (a -> SymBool) ->
    [a] ->
    u ([a], [a])
  mrgSpan _ [] = mrgPure ([], [])
  mrgSpan p xs@(x : xs') =
    mrgIf (p x) (mrgFmap (first (x :)) $ mrgSpan p xs') (mrgPure ([], xs))
  
  -- | Symbolic version of 'Data.List.break', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- Can generate O(n) cases and O(n) sized branch constraints.
  mrgBreak ::
    (Applicative u, SymBranching u, Mergeable a) =>
    (a -> SymBool) ->
    [a] ->
    u ([a], [a])
  mrgBreak p = mrgSpan (symNot . p)
  
  -- | Symbolic version of 'Data.List.stripPrefix', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- Generate O(1) cases and O(len(prefix)) sized branch constraints.
  mrgStripPrefix ::
    (Applicative u, SymBranching u, Mergeable a, SymEq a) =>
    [a] ->
    [a] ->
    u (Maybe [a])
  mrgStripPrefix [] ys = mrgPure $ Just ys
  mrgStripPrefix (x : xs) (y : ys) =
    mrgIf (x .== y) (mrgStripPrefix xs ys) (mrgPure Nothing)
  mrgStripPrefix _ _ = mrgPure Nothing
  
  -- | Symbolic version of 'Data.List.group', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- This function can be very inefficient on large symbolic lists and generate
  -- O(2^n) cases. Use with caution.
  mrgGroup ::
    (MonadUnion u, Mergeable a, SymEq a) =>
    [a] ->
    u [[a]]
  mrgGroup = mrgGroupBy (.==)
  
  -- | Symbolic version of 'Data.List.isPrefixOf'.
  --
  -- Generate O(len(prefix)) sized constraints.
  symIsPrefixOf :: (SymEq a) => [a] -> [a] -> SymBool
  symIsPrefixOf [] _ = con True
  symIsPrefixOf _ [] = con False
  symIsPrefixOf (x : xs) (y : ys) =
    x .== y .&& symIsPrefixOf xs ys
  
  -- | Symbolic version of 'Data.List.isSuffixOf'.
  --
  -- Generate O(len(suffix)) sized constraints.
  symIsSuffixOf :: (SymEq a) => [a] -> [a] -> SymBool
  symIsSuffixOf ns hs = symIsPrefixOf (reverse ns) (reverse hs)
  
  -- | Symbolic version of 'Data.List.isInfixOf'.
  --
  -- Generate O(len(haystack) * len(needle)) sized constraints.
  symIsInfixOf :: (SymEq a) => [a] -> [a] -> SymBool
  symIsInfixOf needle haystack = symAny (symIsPrefixOf needle) (tails haystack)
  
  -- | Symbolic version of 'Data.List.isSubsequenceOf'.
  --
  -- Generate O(len(haystack) * len(needle)) sized constraints.
  symIsSubsequenceOf :: (SymEq a) => [a] -> [a] -> SymBool
  symIsSubsequenceOf [] _ = con True
  symIsSubsequenceOf _ [] = con False
  symIsSubsequenceOf a@(x : a') (y : b) =
    symIte (x .== y) (symIsSubsequenceOf a' b) (symIsSubsequenceOf a b)
  
  -- | Symbolic version of 'Data.List.lookup', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- Can generate O(n) cases and O(n) sized branch constraints.
  mrgLookup ::
    forall a b u.
    (Applicative u, SymBranching u, Mergeable b, SymEq a) =>
    a ->
    [(a, b)] ->
    u (Maybe b)
  mrgLookup _ [] = mrgPure Nothing
  mrgLookup key l =
    mrgIf (symAll (key ./=) (fst <$> l)) (mrgPure Nothing) $
      mrgLookup' l
    where
      mrgLookup' :: [(a, b)] -> u (Maybe b)
      mrgLookup' [] = error "mrgLookup: impossible"
      mrgLookup' [(_, y)] = mrgPure $ Just y
      mrgLookup' ((x, y) : xys) =
        mrgIf (key .== x) (mrgPure $ Just y) (mrgLookup' xys)
  
  -- | Symbolic version of 'Data.List.filter', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- This function can be very inefficient on large symbolic lists and generate
  -- O(2^n) cases. Use with caution.
  mrgFilter ::
    (Applicative u, SymBranching u, Mergeable a) =>
    (a -> SymBool) ->
    [a] ->
    u [a]
  mrgFilter _ [] = mrgPure []
  mrgFilter p (x : xs) =
    mrgIf (p x) (mrgFmap (x :) $ mrgFilter p xs) (mrgFilter p xs)
  
  -- | Symbolic version of 'Data.List.partition', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- This function can be very inefficient on large symbolic lists and generate
  -- O(2^n) cases. Use with caution.
  mrgPartition ::
    forall u a.
    (Applicative u, SymBranching u, Mergeable a) =>
    (a -> SymBool) ->
    [a] ->
    u ([a], [a])
  mrgPartition _ [] = mrgPure ([], [])
  mrgPartition p (x : xs) =
    mrgIf
      (p x)
      (mrgFmap (first (x :)) partitioned)
      (mrgFmap (second (x :)) partitioned)
    where
      partitioned :: u ([a], [a])
      partitioned = mrgPartition p xs
  
  -- | Symbolic version of 'Data.List.!?', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- Can generate O(1) cases and O(n) sized branch constraints.
  (.!?) ::
    forall a uf int.
    ( MonadUnion uf,
      Mergeable a,
      Num int,
      SymEq int
    ) =>
    [a] ->
    int ->
    uf (Maybe a)
  l .!? p = go l p 0
    where
      go :: [a] -> int -> int -> uf (Maybe a)
      go [] _ _ = mrgReturn Nothing
      go (x : xs) p1 i = mrgIf (p1 .== i) (mrgReturn $ Just x) (go xs p1 $ i + 1)
  
  -- | Symbolic version of 'Data.List.elemIndex', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- Can generate O(n) cases (or O(1) if int is merged), and O(n^2) sized
  -- constraints.
  mrgElemIndex ::
    (MonadUnion u, Mergeable int, SymEq a, Num int) =>
    a ->
    [a] ->
    u (Maybe int)
  mrgElemIndex x = mrgFindIndex (x .==)
  
  -- | Symbolic version of 'Data.List.elemIndices', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- Can generate O(n) cases, and O(n^3) sized constraints.
  mrgElemIndices ::
    (MonadUnion u, Mergeable int, SymEq a, Num int) =>
    a ->
    [a] ->
    u [int]
  mrgElemIndices x = mrgFindIndices (x .==)
  
  -- | Symbolic version of 'Data.List.findIndex', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- Can generate O(n) cases (or O(1) if int is merged), and O(n^2) sized
  -- constraints, assuming the predicate only generates O(1) constraints.
  mrgFindIndex ::
    (Applicative u, SymBranching u, Mergeable int, SymEq a, Num int) =>
    (a -> SymBool) ->
    [a] ->
    u (Maybe int)
  mrgFindIndex p l = mrgFmap listToMaybe $ mrgFindIndices p l
  
  -- | Symbolic version of 'Data.List.findIndices', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- Can generate O(n) cases, and O(n^3) sized constraints, assuming the predicate
  -- only generates O(1) constraints.
  mrgFindIndices ::
    forall u a int.
    (Applicative u, SymBranching u, Mergeable int, SymEq a, Num int) =>
    (a -> SymBool) ->
    [a] ->
    u [int]
  mrgFindIndices p xs = go $ zip xs $ fromIntegral <$> [0 ..]
    where
      go :: [(a, int)] -> u [int]
      go [] = mrgPure []
      go ((x, y) : xys) = mrgIf (p x) (mrgFmap (y :) $ go xys) (go xys)
  
  -- | Symbolic version of 'Data.List.nub', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- Can generate O(n) cases, and O(n^3) sized constraints.
  mrgNub ::
    (Applicative u, SymBranching u, Mergeable a, SymEq a) =>
    [a] ->
    u [a]
  mrgNub = mrgNubBy (.==)
  
  -- | Symbolic version of 'Data.List.delete', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- Can generate O(n) cases, and O(n^2) sized constraints.
  mrgDelete ::
    (Applicative u, SymBranching u, Mergeable a, SymEq a) =>
    a ->
    [a] ->
    u [a]
  mrgDelete = mrgDeleteBy (.==)
  
  -- | Symbolic version of 'Data.List.\\', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- Can generate O(len(lhs)) cases, and O(len(lhs)^2 * len(rhs)) sized
  -- constraints.
  (.\\) ::
    (MonadUnion u, Mergeable a, SymEq a) =>
    [a] ->
    [a] ->
    u [a]
  (.\\) = mrgDeleteFirstsBy (.==)
  
  -- | Symbolic version of 'Data.List.union', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- Can generate O(len(rhs)) cases, and O(len(rhs)^5 * len(lhs)) sized
  -- constraints.
  --
  -- Should be improvable.
  mrgUnion ::
    (MonadUnion u, Mergeable a, SymEq a) =>
    [a] ->
    [a] ->
    u [a]
  mrgUnion = mrgUnionBy (.==)
  
  -- | Symbolic version of 'Data.List.intersect', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- Can generate O(len(rhs)) cases, and O(len(lhs) * len(rhs)) constraints.
  mrgIntersect ::
    (MonadUnion u, Mergeable a, SymEq a) =>
    [a] ->
    [a] ->
    u [a]
  mrgIntersect = mrgIntersectBy (.==)
  
  -- | Symbolic version of 'Data.List.nubBy', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- Can generate O(n) cases, and O(n^3) sized constraints, assuming the predicate
  -- only generates O(1) constraints.
  mrgNubBy ::
    forall a u.
    (Applicative u, SymBranching u, Mergeable a) =>
    (a -> a -> SymBool) ->
    [a] ->
    u [a]
  mrgNubBy eq l = mrgNubBy' l []
    where
      mrgNubBy' :: [a] -> [a] -> u [a]
      mrgNubBy' [] _ = mrgPure []
      mrgNubBy' (y : ys) xs =
        mrgIf
          (mrgElemBy y xs)
          (mrgNubBy' ys xs)
          (mrgFmap (y :) $ mrgNubBy' ys (y : xs))
      mrgElemBy _ [] = con False
      mrgElemBy x (y : ys) = eq x y .|| mrgElemBy x ys
  
  -- | Symbolic version of 'Data.List.deleteBy', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- Can generate O(n) cases, and O(n^2) sized constraints, assuming the predicate
  -- only generates O(1) constraints.
  mrgDeleteBy ::
    (Applicative u, SymBranching u, Mergeable a) =>
    (a -> a -> SymBool) ->
    a ->
    [a] ->
    u [a]
  mrgDeleteBy _ _ [] = mrgPure []
  mrgDeleteBy eq x (y : ys) =
    mrgIf (eq x y) (mrgPure ys) (mrgFmap (y :) $ mrgDeleteBy eq x ys)
  
  -- | Symbolic version of 'Data.List.deleteFirstsBy', the result would be merged
  -- and propagate the mergeable knowledge.
  --
  -- Can generate O(len(lhs)) cases, and O(len(lhs)^2 * len(rhs)) sized
  -- constraints, assuming the predicate only generates O(1) constraints.
  mrgDeleteFirstsBy ::
    (MonadUnion u, Mergeable a) =>
    (a -> a -> SymBool) ->
    [a] ->
    [a] ->
    u [a]
  mrgDeleteFirstsBy eq = mrgFoldlM (flip $ mrgDeleteBy eq)
  
  -- | Symbolic version of 'Data.List.unionBy', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- Can generate O(len(rhs)) cases, and O(len(rhs)^5 * len(lhs)) sized
  -- constraints, assuming the predicate only generates O(1) constraints.
  --
  -- Should be improvable.
  mrgUnionBy ::
    (MonadUnion u, Mergeable a) =>
    (a -> a -> SymBool) ->
    [a] ->
    [a] ->
    u [a]
  mrgUnionBy eq xs ys =
    mrgFmap (xs ++) $
      (mrgNubBy eq ys)
        >>= \nubbed -> mrgFoldlM (flip $ mrgDeleteBy eq) nubbed xs
  
  -- | Symbolic version of 'Data.List.intersectBy', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- Can generate O(len(rhs)) cases, and O(len(lhs) * len(rhs)) constraints,
  -- assuming the predicate only generates O(1) constraints.
  mrgIntersectBy ::
    (MonadUnion u, Mergeable a) =>
    (a -> a -> SymBool) ->
    [a] ->
    [a] ->
    u [a]
  mrgIntersectBy _ [] _ = mrgPure []
  mrgIntersectBy _ _ [] = mrgPure []
  mrgIntersectBy eq xs ys = do
    tl <- mrgIntersectBy eq (tail xs) ys
    mrgIf (symAny (eq (head xs)) ys) (mrgReturn $ head xs : tl) (mrgPure tl)
  
  -- | This function can be very inefficient on large symbolic lists and generate
  -- O(2^n) cases. Use with caution.
  mrgGroupBy ::
    (MonadUnion u, Mergeable a) =>
    (a -> a -> SymBool) ->
    [a] ->
    u [[a]]
  mrgGroupBy _ [] = mrgPure []
  mrgGroupBy eq (x : xs) = do
    (ys, zs) <- mrgSpan (eq x) xs
    tl <- mrgGroupBy eq zs
    mrgReturn $ (x : ys) : tl
  
  -- | Symbolic version of 'Data.List.insert', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- Can generate 1 case, and O(n^2) sized constraints.
  mrgInsert ::
    (MonadUnion m, Mergeable a, SymOrd a) =>
    a ->
    [a] ->
    m [a]
  mrgInsert x [] = mrgPure [x]
  mrgInsert x ys@(y : ys') =
    mrgIf (x .<= y) (mrgReturn $ x : ys) (mrgFmap (y :) $ mrgInsert x ys')
  
  -- | Symbolic version of 'Data.List.insertBy', the result would be merged and
  -- propagate the mergeable knowledge.
  --
  -- Can generate 1 case, and O(n^2) sized constraints, assuming the ordering
  -- function only generates O(1) constraints.
  mrgInsertBy ::
    (MonadUnion m, Mergeable a) =>
    (a -> a -> Union Ordering) ->
    a ->
    [a] ->
    m [a]
  mrgInsertBy _ x [] = mrgPure [x]
  mrgInsertBy cmp x ys@(y : ys') = do
    r <- liftUnion $ cmp x y
    case r of
      GT -> mrgFmap (y :) $ mrgInsertBy cmp x ys'
      _ -> mrgReturn $ x : ys
  