Flatten nested disjunction

Author: simonmar

glean/db/Glean/Query/Flatten.hs

@@ -506,7 +506,7 @@ flattenStmts (Statements s) = map Floating (reverse s)
 
 disjunction :: [FlatStatementGroup] -> FlatStatement
 disjunction [FlatStatementGroup [x]] = unOrdered x
-disjunction groups = FlatDisjunction groups
+disjunction groups = flatDisjunction groups
 
 mkGroup :: [Statements] -> [Statements] -> FlatStatementGroup
 mkGroup ords floats =

glean/db/Glean/Query/Flatten/Types.hs

@@ -21,6 +21,7 @@ module Glean.Query.Flatten.Types
   , Ordered(..)
   , unOrdered
   , grouping
+  , flatDisjunction
   , mkStatementGroup
   , singletonGroup
   , boundVars
@@ -129,6 +130,18 @@ grouping :: FlatStatementGroup -> FlatStatement
 grouping (FlatStatementGroup [one]) = unOrdered one
 grouping group = FlatDisjunction [group]
 
+-- | Flatten nested FlatDisjunction
+flatDisjunction :: [FlatStatementGroup] -> FlatStatement
+flatDisjunction groups =
+  FlatDisjunction
+    [ group
+    | g <- groups
+    , group <- case g of
+        FlatStatementGroup [Ordered (FlatDisjunction gs)] -> gs
+        FlatStatementGroup [Floating (FlatDisjunction gs)] -> gs
+        _ -> [g]
+    ]
+
 -- | Smart constructor for a FlatStatementGroup, as with "grouping"
 -- this flattens unnecessary nesting.
 mkStatementGroup :: [Ordered FlatStatement] -> FlatStatementGroup

glean/db/Glean/Query/Opt.hs

@@ -248,7 +248,7 @@ instance Apply FlatStatement where
       [] -> case stmtss' of
         [] -> return (FlatDisjunction [])
         (ss : _) -> return (grouping ss)
-      some -> return (FlatDisjunction some)
+      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
@@ -731,7 +731,7 @@ filterStmt stmt = case stmt of
   FlatNegation stmts -> FlatNegation <$> filterGroupEnclosed stmts
   FlatDisjunction [stmts] -> grouping <$> filterGroup stmts
   FlatDisjunction stmtss ->
-    FlatDisjunction <$> mapM filterGroupEnclosed stmtss
+    flatDisjunction <$> mapM filterGroupEnclosed stmtss
   FlatConditional cond then_ else_ -> do
     (cond', then') <- encloseSeen $ do
       cond' <- filterGroup cond