{-
  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.Query.Opt
  ( optimise
  ) where

import Control.Monad
import Control.Monad.Except
import Control.Monad.State.Strict
import qualified Data.ByteString as B
import Data.Foldable (toList)
import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap
import qualified Data.IntSet as IntSet
import Data.Text (Text)

import qualified Glean.Angle.Types as Type
import Glean.Angle.Types (FieldDef_(..))
import Glean.Query.Flatten.Types
import Glean.Query.Codegen.Types
import Glean.Query.Vars
import Glean.RTS.Term
import Glean.RTS.Types

{-
Note [Query Simplification]

This pass performs obvious optimisations on a query as follows:

- Perform unification on all statements of the form P = Q
- Apply the substitution resulting from unification

So for example if we have

  P = cxx1.Name S
  S = "abc"

Unification will produce the substitution [S->"abc"]; applying it will produce

  P = cxx1.Name "abc"
  "abc" = "abc"

and finally in the Reorder pass we will identify the statement "abc" =
"abc" as redundant and eliminate it.


Optimising Derived Predicates
-----------------------------

This optimisation is particularly useful when expanding derived
predicates. For example, given

   predicate RevStringPair : { a : string, b : string }
     { A, B } where StringPair { B, A }

and a query

   RevStringPair { _, "a" }

after expansion we will have

   X where StringPair { B, A }; X = { A, B }; X = { _, "a" }

the StringPair query is not efficient here; we could do much better if
we knew that B = "a".

Unification will produce the substitution

  [ X -> { A, B }, B -> "a" ]

and after applying the substitution and simplification we will have

   { A, "a" } where StringPair { "a", A }

which is exactly what we wanted.


Dealing with Choice
-------------------

Choice is the main source of complexity here.

For example, suppose we have

   Y where
     cxx1.Name X;
     Y = (1 where X = "a") | (2 where X = "b")

Clearly there must be some "local" unification, because the
substitution X = "a" only applies inside (1 where X = "a").  But just
applying the local substitution [X->"a"] inside (1 where X = "a")
can't be right, because we would get

   Y = (1 where "a" = "a") | ...

The problem is that we bound X during unification and then discarded
that part of the substitution. We must never discard a substitution
when it has not been applied to all instances of a variable - and yet
we cannot let the substitition for X escape from the branch of this
choice, because it isn't valid outside. This is a contradiction,
so the solution is not to unify X in the first place. But how do we
know we shouldn't unify X? Because it is visible *outside* the scope
of the local substitution.

So the rule for unificaiton of (A | B) is:

- Unification is local to each branch. We have a substitution S
  arising from the context and two separate substitutions Sa and Sb
  arising from unification inside A and B respectively. The
  subsititions we apply to A and B are S+Sa and S+Sb respectively.

- We must *not* unify any variables that are visible outside. In other
  words, Sa and Sb can only bind variables that are local to A and B
  respectively.


Algorithm
---------

Unification of a complete query proceeds as follows:

0. Start with the empty substitution S = {}
1. Unify everywhere except (A | B) subterms, extending S
2. Apply S everywhere except (A | B) subterms
3. For each (A | B) subterm, repeat from step 1 for A and B separately,
   starting with current substitition S

-}

optimise :: FlattenedQuery -> Except Text FlattenedQuery
optimise query@QueryWithInfo{..} = do
  let
    state = OptState
      { optNextVar = qiNumVars
      , optSubst = IntMap.empty
      , optGenerators = IntMap.empty
      , optCurrentScope = IntSet.empty
      , optOuterScope = IntSet.empty
      , optSeen = IntMap.empty
      }
  (optimised, state') <- flip runStateT state $
    optimiseQuery =<< freshWildQuery qiQuery
    -- unification relies on wildcards being replaced by variables
    --  1. so that when we substitute an expression in multiple places
    --     we can unify the wildcards across the different instantiations.
    --  2. we want to elimiante duplicate "X = generator" statements
    --     by unifying the generators and then spotting identical
    --     statements. But those will only end up identical if we
    --     replace the wildcards with variables before unifying and
    --     substituting.
    -- freshWildQuery also replaces MatchBind with MatchVar, so we don't
    -- have spurious differences between terms that are really identical.
  return query { qiQuery = optimised, qiNumVars = optNextVar state' }

optimiseQuery :: FlatQuery -> U FlatQuery
optimiseQuery query@(FlatQuery key maybeVal stmts) = do
  -- determine variables visible outside of any nested choice:
  modify $ \s -> s { optCurrentScope = queryScope query }
  stmts' <- optStmts stmts
  FlatQuery
    <$> apply key
    <*> mapM apply maybeVal
    <*> pure stmts'

-- | Apply the current substutition
class Apply t where
  apply :: t -> U t

instance Apply Generator where
  apply (FactGenerator pid key val range) =
    FactGenerator pid <$> apply key <*> apply val <*> pure range
  apply (TermGenerator x) = TermGenerator <$> apply x
  apply (DerivedFactGenerator pid key val) =
    DerivedFactGenerator pid <$> apply key <*> apply val
  apply (ArrayElementGenerator ty arr) =
    ArrayElementGenerator ty <$> apply arr
  apply (SetElementGenerator ty arr) =
    SetElementGenerator ty <$> apply arr
  apply (PrimCall op args ty) =
    PrimCall op <$> mapM apply args <*> pure ty

{- Note [unification failure]

Unification failure can easily arise as a result of choice. For
example:

(0|X = 0|"abc") | (0|X = 1|"def")

The unification of (0|X) with (1|"def") fails, and we would like to
exploit that to simplify the whole statement to

(0|X) = (0|"abc")

This works as follows:

* unify returns a Bool to indicate whether unification succeeded

* unifyStmt replaces a statement with falseStmt if unification of the
  LHS and RHS fails

* In a disjunctions, we filter out any alternatives that are
  definitely empty, where empty is defined as "any statement is
  obviously false".

Note: we have to be careful about getting rid of statements entirely,
even if a statement is provably false, because deleting a statement
can leave unbound variables behind.  See optStmts and apply below.

-}

instance Apply FlatStatementGroup where
  apply (FlatStatementGroup stmts) =
    mkStatementGroup <$> mapM (mapM apply) stmts

instance Apply FlatStatement where
  apply (FlatStatement ty lhs rhs) = do
    -- these were already unified by 'unify'
    FlatStatement ty <$> apply lhs <*> apply rhs
  apply (FlatAllStatement v e g) = do
    t <- applyVar v
    -- it must be a variable, because no other pattern matches a set.
    v' <- case t of
      Ref (MatchVar x) -> return x
      Ref (MatchBind x) -> return x
      _ -> error "apply: FlatAllStatement"
    enclose (stmtGroupScope g . termScope e) $ do
      g' <- optStmts g
      FlatAllStatement v' <$> apply e <*> pure g'
  apply (FlatNegation stmts) = do
    -- assumptions arising inside the negation are not true outside of it.
    stmts' <- optStmtsEnclosed stmts
    return (FlatNegation  stmts')
  apply (FlatDisjunction [stmts]) = do
    stmts' <- apply stmts
    return (grouping stmts')
  apply (FlatDisjunction stmtss) = do
    stmtss' <- mapM optStmtsEnclosed stmtss
    -- we can remove alternatives that are known to be false, but if
    -- ALL alternatives are false, then we must leave behind just one
    -- alternative so that we don't leave unbound variables.
    case filter (not . isFalseGroups) stmtss' of
      [] -> case stmtss' of
        [] -> return (FlatDisjunction [])
        (ss : _) -> return (grouping ss)
      some -> return (FlatDisjunction some)
  apply (FlatConditional cond then_ else_) = do
    -- like disjunctions, assumptions arising from the conditional statements
    -- are not true outside of it. However, those arising from the condition
    -- are true in the 'then' case.
    (cond', then') <-
      enclose (stmtGroupScope cond) $ do
        cond' <- optStmts cond
        then' <- optStmtsEnclosed then_
        return (cond', then')
    else' <- optStmtsEnclosed else_

    return $ if isFalseGroups cond'
      then grouping else'
      else FlatConditional cond' then' else'

optStmtsEnclosed :: FlatStatementGroup -> U FlatStatementGroup
optStmtsEnclosed stmts = enclose (stmtGroupScope stmts) $ optStmts stmts

-- If a sequence of statements is found to be false, then we place
-- a falseStmt sentinel at the beginning. We don't actually remove
-- any statements at this stage, because we have to be careful not to
-- leave any variables unbound.
optStmts :: FlatStatementGroup -> U FlatStatementGroup
optStmts (FlatStatementGroup ord) = do
  notFalse <- and <$> mapM unifyOrdStmt ord
  ord' <- concatMap (mapM expandStmt) <$> mapM (mapM apply) ord
  ord'' <- filterOrdStmts ord'
  -- unify may fail, but apply may also leave behind a false statement:
  if notFalse && not (any isFalseOrdStmt ord'')
    then return (mkStatementGroup ord'')
    else return (mkStatementGroup (Ordered falseStmt : ord'' ))

-- Look for the sentinel left by optStmts
isFalseGroups :: FlatStatementGroup -> Bool
isFalseGroups (FlatStatementGroup (s : _)) = isFalseOrdStmt s
isFalseGroups (FlatStatementGroup []) = False

isFalseOrdStmt :: Ordered FlatStatement -> Bool
isFalseOrdStmt = isFalseStmt . unOrdered

isFalseStmt :: FlatStatement -> Bool
isFalseStmt (FlatNegation (FlatStatementGroup [])) = True
isFalseStmt (FlatDisjunction []) = True
isFalseStmt (FlatDisjunction [FlatStatementGroup ord]) =
  any isFalseOrdStmt ord
isFalseStmt _ = False

instance Apply Pat where
  apply (Tuple xs) = Tuple <$> mapM apply xs
  apply (Array xs) = Array <$> mapM apply xs
  apply (Alt n x) = Alt n <$> apply x
  apply (Ref (MatchVar v)) = applyVar v
  apply (Ref (MatchBind v)) = applyVar v
  apply (Ref (MatchAnd x y)) = do
    x' <- apply x
    y' <- apply y
    return (Ref (MatchAnd x' y'))
  apply (Ref (MatchPrefix str x)) = Ref . MatchPrefix str <$> apply x
  apply (Ref (MatchArrayPrefix ty pre all)) =
    Ref <$> (MatchArrayPrefix ty <$> mapM apply pre <*> apply all)
  apply t@Byte{} = return t
  apply t@ByteArray{} = return t
  apply t@Nat{} = return t
  apply t@String{} = return t
  apply t@(Ref MatchExt{}) = return t
  apply t@(Ref MatchFid{}) = return t
  apply t@(Ref MatchNever{}) = return t
  apply t@(Ref MatchWild{}) = return t

applyVar :: Var -> U Pat
applyVar var@(Var _ v _) = do
  subst <- gets optSubst
  case IntMap.lookup v subst of
    Just pat -> apply pat
      -- also apply the substitition to the result, because we don't
      -- keep the substitution already in ground form. This is
      -- potentially inefficient and there are better ways to do it
      -- (typically with mutable variables), but it's probably not
      -- important enough to fix yet.
    _otherwise -> return (Ref (MatchVar var))

-- | For running unification on one alternative of nested choice
-- or a negated query.
--
-- Does not add substitutions or new variables to the parent scope.
enclose :: (VarSet -> VarSet) -> U a -> U a
enclose innerScope u = do
  state0 <- get
  -- set the outer scope to be the current scope
  let scope = innerScope (optCurrentScope state0)
  modify $ \s ->
    s { optCurrentScope = scope, optOuterScope = optCurrentScope state0 }
  a <- u
  modify $ \s ->
    s { optSubst = optSubst state0
      , optGenerators = optGenerators state0
      , optCurrentScope = optCurrentScope state0
      , optOuterScope = optOuterScope state0
      , optSeen = optSeen state0 }
  return a

unifyOrdStmt :: Ordered FlatStatement -> U Bool
unifyOrdStmt = unifyStmt . unOrdered

unifyStmt :: FlatStatement -> U Bool
unifyStmt (FlatStatement _ lhs rhs)
  | neverMatches lhs = return False
  | TermGenerator rhs' <- rhs = unify lhs rhs'
  | otherwise = do
    r <- case isVar lhs of
      Just v -> extend v rhs
      Nothing -> return True
    case rhs of
      FactGenerator _ key val _ ->
        return $ r && not (neverMatches key || neverMatches val)
      _ -> return r
unifyStmt FlatAllStatement{} = return True
unifyStmt FlatNegation{} = return True
  -- ignore negations for now. We will recurse into it later
unifyStmt (FlatDisjunction [FlatStatementGroup ord]) = do
  as <- mapM unifyOrdStmt ord
  return $ and as
  -- singleton FlatDisjunction is used for grouping, we must retain
  -- it, but not treat it as a disjunction.
unifyStmt FlatDisjunction{} = return True
  -- ignore a disjunction for now. We will recurse into it in 'apply'
unifyStmt FlatConditional{} = return True
  -- ignore conditions for now. We will recurse into it in 'apply'

neverMatches :: Pat -> Bool
neverMatches = \case
  Byte _ -> False
  Nat _ -> False
  Array terms -> any neverMatches terms
  ByteArray _ -> False
  String _ -> False
  Tuple terms -> any neverMatches terms
  Alt _ term -> neverMatches term
  Ref match -> case match of
    MatchWild _ -> False
    MatchNever _ -> True
    MatchFid _ -> False
    MatchBind _ -> False
    MatchVar _ -> False
    MatchAnd left right -> neverMatches left || neverMatches right
    MatchPrefix _ term -> neverMatches term
    MatchArrayPrefix _ty pre all -> any neverMatches pre || neverMatches all
    MatchExt () -> False

unifyGen :: Generator -> Generator -> U Bool
unifyGen (TermGenerator a) (TermGenerator b) = a `unify` b
unifyGen (FactGenerator _ k1 v1 _) (FactGenerator _ k2 v2 _) =
  (&&) <$> unify k1 k2 <*> unify v1 v2
unifyGen (DerivedFactGenerator _ k1 v1) (DerivedFactGenerator _ k2 v2) =
  (&&) <$> unify k1 k2 <*> unify v1 v2
unifyGen _ _ = return True

-- | Unify two patterns, extending the substitution. Note that this is
-- not complete unification: we aren't guaranteeing to produce a
-- unifier. We might refrain from unifying in some cases (see the
-- comment about or-patterns at the top of the module).
unify :: Pat -> Pat -> U Bool
unify (Byte x) (Byte y) = return (x == y)
unify (Nat x) (Nat y) = return (x == y)
unify (ByteArray x) (ByteArray y) = return (x == y)
unify (String x) (String y) = return (x == y)
unify (Ref (MatchFid x)) (Ref (MatchFid y)) = return (x == y)
unify (Ref MatchWild{}) _ = return True
unify _ (Ref MatchWild{}) = return True
unify (Ref MatchNever{}) _ = return False
unify _ (Ref MatchNever{}) = return False
unify (Ref (MatchVar v)) pat = extend v (TermGenerator pat)
unify pat (Ref (MatchVar v)) = extend v (TermGenerator pat)
unify (Ref (MatchBind v)) pat = extend v (TermGenerator pat)
unify pat (Ref (MatchBind v)) = extend v (TermGenerator pat)
unify (Tuple xs) (Tuple ys) = and <$> zipWithM unify xs ys
unify (Array xs) (Array ys)
  | length xs == length ys = and <$> zipWithM unify xs ys
  | otherwise = return False
unify (Alt n x) (Alt m y)
  | n == m = unify x y
  | otherwise = return False
unify (Ref (MatchPrefix s x)) (Ref (MatchPrefix t y))
  | s == t = unify x y
  | otherwise = return True -- might still match
unify (Ref (MatchPrefix s x)) (String y)
  | Just r <- B.stripPrefix s y = unify x (String r)
unify (String x) (Ref (MatchPrefix s y))
  | Just r <- B.stripPrefix s x = unify y (String r)
unify (Ref (MatchAnd x y)) z = (&&) <$> unify x z <*> unify y z
unify z (Ref (MatchAnd x y)) = (&&) <$> unify z x <*> unify z y
unify (Ref (MatchArrayPrefix _ pre all)) (Ref (MatchArrayPrefix _ pre' all')) = do
  b <- unify all all'
  bs <- zipWithM unify pre pre'
  return (and (b:bs))
unify (Ref (MatchArrayPrefix _ pre all)) (Array xs)
  | length pre <= length xs = do
    b <- unify all (Array xs)
    bs <- zipWithM unify pre xs
    return (and (b:bs))
unify a@Array{} b@(Ref (MatchArrayPrefix{})) = unify b a
unify _ _ = return False

extend :: Var -> Generator -> U Bool
-- These two generators don't tell us anything:
extend _ ArrayElementGenerator{} = return True
extend _ SetElementGenerator{} = return True
extend _ PrimCall{} = return True
-- TermGenerators get added to the substitution:
extend var (TermGenerator t) = do
  state <- get
  t <- apply t  -- ensure we don't create a recursive substitution
  let
    bind v gen = do
      modify $ \s -> s { optSubst = IntMap.insert (varId v) gen (optSubst s) }
      return True

    hasSubstitution v = IntMap.lookup (varId v) (optSubst state)
    comesFromOuterScope v = IntSet.member (varId v) (optOuterScope state)

    -- refuse to bind a variable if it is visible in the outer scope
    -- (see comment at the top of this module). If this happens, we
    -- try the unification the other way around instead.
    check v t _ | Just v' <- isVar t, v == v' = return True
    check v t orElse
      | Just t' <- hasSubstitution v = unify t t'
      | comesFromOuterScope v = orElse
      | otherwise = bind v t

  check var t $
    case isVar t of
      Just y -> check y (Ref (MatchVar var)) (return True)
      _ -> extendGen var (TermGenerator t)
-- Other kinds of generators: remember them in optGenerators
extend var gen = extendGen var gen

extendGen :: Var -> Generator -> U Bool
extendGen var gen = do
  state <- get
  gen <- apply gen
  pat <- applyVar var
  case isVar pat of
    Nothing ->
      -- it can be a fact ID. What should we do in that case? (TODO)
      return True
    Just var' ->
      case IntMap.lookup (varId var') (optGenerators state) of
        Just gen' -> unifyGen gen gen'
        Nothing -> do
          modify $ \s -> s { optGenerators =
            IntMap.insert (varId var') gen (optGenerators s) }
          return True

data OptState = OptState
  { optNextVar :: !Int
    -- ^ for making fresh variables
  , optSubst :: Subst
    -- ^ current substitution
  , optGenerators :: IntMap Generator
    -- ^ knowledge about which variables are bound by generators. Only
    -- contains FactGenerators and DerivedFactGenerators currently.
  , optCurrentScope :: VarSet
    -- ^ all variables visible in the current scope, excluding nested
    -- (A | B) subterms
  , optOuterScope :: VarSet
    -- ^ all variables visible in the enclosing scope, when we are
    -- unifying inside a nested (A | B) subterm.  These are the
    -- variables we must not unify (see comment at the top of the
    -- module).
  , optSeen :: IntMap Generator
    -- ^ statements we have seen before. See Note [identical statements]
  }

instance Monad m => Fresh (StateT OptState m) where
  peek = gets optNextVar
  alloc = do
    state@OptState{..} <- get
    put state{ optNextVar = optNextVar + 1 }
    return optNextVar

type U a = StateT OptState (Except Text) a

type Subst = IntMap Pat

-- | Traverse the query, excluding nested (A | B), !A, and "if" subterms,
-- collecting the visible variables.  These are the variables that we must
-- not unify *inside* an (A | B), !A, or "if" subterm, because they are
-- visible outside it. Variables that are local to one of these subterms
-- may be safely unified.
queryScope :: FlatQuery -> VarSet
queryScope (FlatQuery key maybeVal (FlatStatementGroup ord)) =
  foldr termScope (foldr ordStmtScope s ord) maybeVal
  where
    s = termScope key IntSet.empty

stmtGroupScope :: FlatStatementGroup -> VarSet -> VarSet
stmtGroupScope (FlatStatementGroup g) r = foldr ordStmtScope r g

ordStmtScope :: Ordered FlatStatement -> VarSet -> VarSet
ordStmtScope = stmtScope . unOrdered

stmtScope :: FlatStatement -> VarSet -> VarSet
stmtScope (FlatStatement _ lhs rhs) r = termScope lhs (genScope rhs r)
stmtScope (FlatAllStatement v _ (FlatStatementGroup _)) r =
  addToCurrentScope v $! r
stmtScope (FlatNegation _) r = r
stmtScope (FlatDisjunction [FlatStatementGroup ord]) r =
  foldr ordStmtScope r ord
stmtScope FlatDisjunction{} r = r
  -- contents of or-patterns are not part of the "current scope"
stmtScope FlatConditional{} r = r
  -- contents of if-patterns are not part of the "current scope"

genScope :: Generator -> VarSet -> VarSet
genScope (FactGenerator _ key val _) r = termScope key $! termScope val r
genScope (TermGenerator pat) r = termScope pat r
genScope (DerivedFactGenerator _ key val) r = termScope key $! termScope val r
genScope (ArrayElementGenerator _ exp) r = termScope exp r
genScope (SetElementGenerator _ exp) r = termScope exp r
genScope (PrimCall _ args _) r = foldr termScope r args

addToCurrentScope :: Var -> VarSet -> VarSet
addToCurrentScope (Var _ v _) r = IntSet.insert v r

termScope :: Pat -> VarSet -> VarSet
termScope pat r = foldr onMatch r pat
  where
  onMatch :: Match () Var -> VarSet -> VarSet
  onMatch m r = case m of
    MatchPrefix _ rest -> foldr onMatch r rest
    MatchBind v -> addToCurrentScope v r
    MatchVar v -> addToCurrentScope v r
    MatchAnd x y -> foldr onMatch (foldr onMatch r y) x
    MatchArrayPrefix _ty pre all ->
      foldr onMatch (foldr onMatch r all) (foldMap toList pre)
    MatchNever{} -> r
    MatchWild{} -> r
    MatchExt{} -> r
    MatchFid{} -> r


-- | For T = U statements, decompose the statement as far as possible.
-- For example
--
-- >    {A,B} = {C,D}
--
-- decomposes into
--
-- >     A = C
-- >     B = D
--
-- and eliminate any statements that are obviously redundant. That is,
-- when the lhs obviously matches the rhs, and the statement binds no
-- variables.
--
expandGroup :: FlatStatementGroup -> FlatStatementGroup
expandGroup (FlatStatementGroup ord) =
  mkStatementGroup (concatMap (mapM expandStmt) ord)

expandStmt :: FlatStatement -> [FlatStatement]
expandStmt (FlatStatement stmtTy lhs (TermGenerator rhs)) =
  [ FlatStatement ty a (TermGenerator b)
  | (ty,a,b) <- expand stmtTy lhs rhs ]
  where
  expand :: Type -> Pat -> Pat -> [(Type,Pat,Pat)]
  expand ty a b = case (a,b) of
    (_, Ref MatchWild{}) -> []
    (Ref MatchWild{}, _) -> []
    (_, Ref MatchNever{}) -> []
    (Ref MatchNever{}, _) -> []
    (Byte a, Byte b) | a == b -> []
    (Nat a, Nat b) | a == b -> []
    (String a, String b) | a == b -> []
    (Array ts, Array us)
      | Type.ArrayTy eltTy <- derefType ty
      , length ts == length us -> concat (zipWith (expand eltTy) ts us)
    (ByteArray b, ByteArray c) | b == c -> []
    (Tuple ts, Tuple us)
      | Type.RecordTy fields <- derefType ty ->
        concat (zipWith3 expand (map fieldDefType fields) ts us)
    (Alt n x, Alt m y)
      | n == m, Just fields <- sumLike (derefType ty) ->
        expand (fieldDefType (fields !! fromIntegral n)) x y
    (Ref x, Ref y)
      | Just x' <- matchVar x, Just y' <- matchVar y, x' == y' -> []
    (Ref (MatchAnd a b), x) -> expand ty a x ++ expand ty b x
    (x, Ref (MatchAnd a b)) -> expand ty x a ++ expand ty x b
    (Ref (MatchPrefix s a), String b)
      | Just r <- B.stripPrefix s b -> expand ty a (String r)
    (String a, Ref (MatchPrefix s b))
      | Just r <- B.stripPrefix s a -> expand ty (String r) b
    (Ref (MatchPrefix s a), Ref (MatchPrefix t b))
      | t == s -> expand ty a b
    (Ref (MatchFid x), Ref (MatchFid y)) | x == y -> []
    (Ref (MatchArrayPrefix eltTy0 ts all), Array us)
      | length ts <= length us
      , eltTy <- derefType eltTy0
      -> concat (expand ty all (Array us) : zipWith (expand eltTy) ts us)
    (Ref (MatchArrayPrefix eltTy0 ts all), Ref (MatchArrayPrefix _ty us all'))
      | eltTy <- derefType eltTy0
      , length ts == length us
      -> concat (expand ty all all' : zipWith (expand eltTy) ts us)
    _ -> [(ty,a,b)]
expandStmt s@(FlatNegation _) = [s]
  -- the expansion is handled by @apply@ with @optStmts@ because we need to
  -- constrain the scope of the substitutions arising from inside the negation.
expandStmt (FlatDisjunction [stmts]) =
  case expandGroup stmts of
    FlatStatementGroup [] -> []
    xs -> [grouping xs]
  -- non-singleton disjunctions are handled by apply, which will
  -- expand the stmts via optStmts.
expandStmt s = [s]

-- -----------------------------------------------------------------------------
-- Eliminating duplicate statements

{-
Note [identical statements]

We want to combine multiple generators like

  X = pred { _, A }
  X = pred { B, "x" }

this works as follows: first we replace all the wildcards with fresh variables

  X = pred { _1, A }
  X = pred { B, "x" }

Next, when unifying we insert the first statement into the substitution

  { X := pred { _1, A } }

when we see the second statement, we unify its generator with the
first one (see 'extend'), giving

  { _1 := B, A := "x", X := pred { _1, A } }

in 'apply', we apply the substitution to each statement:

  X = pred { B, "x" }
  X = pred { B, "x" }

next, we just elimiante literally identical statements, being careful
that we don't propagate assumptions from disjunctions/negation to the
outside scope, as with unification.

Note that even after unification and substitution we might still have
different generators for a given variable, when some of them occur
inside disjunctions or negation.
-}

notDuplicateStmt :: FlatStatement -> U Bool
notDuplicateStmt (FlatStatement _ lhs rhs)
  | Ref (MatchVar v) <- lhs = dup v rhs
  | Ref (MatchBind v) <- lhs = dup v rhs
  where
  dup (Var _ v _ ) rhs = do
    history <- gets optSeen
    case IntMap.lookup v history of
      Just gen -> return (gen /= rhs)
      Nothing -> do
        modify $ \s -> s { optSeen = IntMap.insert v rhs history }
        return True
notDuplicateStmt _other = return True

encloseSeen :: U a -> U a
encloseSeen inner = do
  state <- get
  a <- inner
  modify $ \s -> s { optSeen = optSeen state }
  return a

filterStmt :: FlatStatement -> U FlatStatement
filterStmt stmt = case stmt of
  FlatStatement{} -> return stmt
  FlatAllStatement v e stmts ->
    FlatAllStatement v e <$> filterGroupEnclosed stmts
  FlatNegation stmts -> FlatNegation <$> filterGroupEnclosed stmts
  FlatDisjunction [stmts] -> grouping <$> filterGroup stmts
  FlatDisjunction stmtss ->
    FlatDisjunction <$> mapM filterGroupEnclosed stmtss
  FlatConditional cond then_ else_ -> do
    (cond', then') <- encloseSeen $ do
      cond' <- filterGroup cond
      then' <- filterGroupEnclosed then_
      return (cond', then')
    else' <- filterGroupEnclosed else_
    return (FlatConditional cond' then' else')

filterGroup :: FlatStatementGroup -> U FlatStatementGroup
filterGroup (FlatStatementGroup ord) =
  mkStatementGroup <$> filterOrdStmts ord

filterOrdStmts :: [Ordered FlatStatement] -> U [Ordered FlatStatement]
filterOrdStmts stmts = do
  filtered <- filterM (notDuplicateStmt . unOrdered) stmts
  mapM (mapM filterStmt) filtered

filterGroupEnclosed :: FlatStatementGroup -> U FlatStatementGroup
filterGroupEnclosed = encloseSeen . filterGroup