Code generation constructor for 'all'

Summary:
Very similar to the flattened representation, so no surprises here.

Upcoming diffs will contain the code generation bits.

Reviewed By: simonmar

Differential Revision: D60139166

fbshipit-source-id: 768e1c4779f2d05ce2b65870b6596b73a0ffcb09

glean/db/Glean/Query/Codegen.hs

@@ -132,6 +132,8 @@ findOutputs q = findOutputsQuery q IntSet.empty
   findOutputsStmt :: CgStatement -> IntSet -> IntSet
   findOutputsStmt (CgStatement lhs gen) r =
      findOutputsGen gen (foldr findOutputsMatch r lhs)
+  findOutputsStmt (CgAllStatement (Var _ var _) expr stmts) r =
+     IntSet.insert var $ findOutputsPat expr $ foldr findOutputsStmt r stmts
   findOutputsStmt (CgNegation stmts) r =
      foldr findOutputsStmt r stmts
   findOutputsStmt (CgDisjunction stmtss) r =
@@ -496,6 +498,8 @@ compileStatements
           matches :: Term (Match e v) -> [Match e v]
           matches p = foldMap pure p
 
+      compile (CgAllStatement (Var _ _var _) _expr _stmts : _rest) =
+        error "Set"
       compile (CgNegation stmts : rest) = mdo
         singleResult (compileStatements qtrans bounds regs stmts vars) $
           jump fail

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

@@ -76,6 +76,7 @@ data CgQuery = CgQuery
 --
 data CgStatement_ var
   = CgStatement (Pat_ var) (Generator_ var)
+  | CgAllStatement var (Expr_ var) [CgStatement_ var]
   | CgNegation [CgStatement_ var]
   | CgDisjunction [[CgStatement_ var]]
   -- ^ For rationale, see Note [why do we have sequential composition?]
@@ -310,6 +311,9 @@ instance Display CgStatement where
   display opts = \case
     CgStatement pat gen ->
       hang 2 $ sep [display opts pat <+> "=", display opts gen]
+    CgAllStatement var expr stmts ->
+      hang 2 $ sep [display opts var <+> "="
+                   ,display opts (CgQuery expr stmts)]
     CgNegation stmts -> "!" <> doStmts stmts
     CgDisjunction stmtss -> sep (punctuate " |" (map doStmts stmtss))
     CgConditional cond then_ else_ -> sep

glean/db/Glean/Query/Incremental.hs

@@ -86,6 +86,7 @@ expandGenerators hasFacts stmts =
       Incremental (generator SeekOnStacked) (generator SeekOnBase)
     CgStatement{} ->
       Incremental Nothing (Just stmt)
+    CgAllStatement{} -> error "Set"
     CgDisjunction stmtss ->
       CgDisjunction <$> expandDisjunction stmtss
     CgNegation{} ->

glean/db/Glean/Query/Reorder.hs

@@ -761,8 +761,7 @@ toCgStatement stmt = case stmt of
     bindVar v
     e' <- fixVars IsExpr e
     stmts' <- mapM toCgStatement stmts
--- TODO: return [CgAllStatement v e' (concat stmts')]
-    error "Set" e' stmts'
+    return [CgAllStatement v e' (concat stmts')]
   FlatNegation stmts -> do
     stmts' <-
       withinNegation $

glean/db/Glean/Query/Transform.hs

@@ -222,6 +222,9 @@ buildTerm output vars term = go term
       outputNatImm (fromIntegral (length vs)) output
       mapM_ go vs
     Tuple fields -> mapM_ go fields
+    Set vs -> do
+      outputNatImm (fromIntegral (length vs)) output
+      mapM_ go vs
     Alt n term -> do outputNatImm n output; go term
     Ref (MatchFid f) -> outputNatImm (fromIntegral (fromFid f)) output
     Ref (MatchPrefix str rest) -> do

glean/db/Glean/Query/Vars.hs

@@ -96,6 +96,9 @@ instance VarsOf CgQuery where
 
 instance VarsOf CgStatement where
   varsOf w (CgStatement lhs gen) r = varsOf w lhs $! varsOf w gen r
+  varsOf w (CgAllStatement (Var _ v _) expr stmts) r =
+    varsOf w expr $! varsOf w stmts $!
+    if w /= VarsUsed then IntSet.insert v r else r
   varsOf w (CgNegation stmts) r = varsOf w stmts r
   varsOf w (CgDisjunction stmtss) r = varsOf w stmtss r
   varsOf w (CgConditional cond then_ else_) r =
@@ -183,6 +186,13 @@ reWildGenerator used gen = case gen of
 reWildStatement :: VarMap -> CgStatement -> CgStatement
 reWildStatement used (CgStatement lhs rhs) =
   CgStatement (reWild used lhs) (reWildGenerator used rhs)
+reWildStatement used s@(CgAllStatement (Var ty n x) expr stmts) =
+  case IntMap.lookup n used of
+    Nothing -> error $ "reWildStatement: " <> show (displayVerbose s)
+    Just new ->
+      CgAllStatement
+        (Var ty new x) (reWild used expr)
+        (map (reWildStatement used) stmts)
 reWildStatement used (CgNegation stmts) =
   CgNegation (map (reWildStatement used) stmts)
 reWildStatement used (CgDisjunction stmtss) =

glean/hs/Glean/RTS/Traverse.hs

@@ -19,12 +19,13 @@ import Glean.RTS.Foreign.Bytecode
 import Glean.RTS.Types
 import Glean.RTS.Bytecode.Code
 import Glean.RTS.Bytecode.Gen.Issue
+import GHC.Stack (HasCallStack)
 
 -- | Type tag for Subroutine
 data CompiledTraversal
 
 traversal
-  :: Register ('Fun '[ 'Word, 'Word ])
+  :: HasCallStack => Register ('Fun '[ 'Word, 'Word ])
   -> Register 'DataPtr
   -> Register 'DataPtr
   -> Type
@@ -72,7 +73,14 @@ traversal callback input inputend ty = go False (repType ty)
             return alt
           end <- label
           return ()
-    go _ (SetRep _elty) = error "Set"
+    go refs (SetRep elty) = when (refs || hasRefs elty) $ local $ \size -> mdo
+      inputNat input inputend size
+      jumpIf0 size end
+      loop <- label
+      go True elty
+      decrAndJumpIfNot0 size loop
+      end <- label
+      return ()
     go _ (PredicateRep (Pid pid)) = local $ \ide -> do
       inputNat input inputend ide
       pidr <- constant (fromIntegral pid)
@@ -83,7 +91,7 @@ traversal callback input inputend ty = go False (repType ty)
 -- contained in it.
 --
 -- NOTE: The subroutine assumes that the clause is type correct.
-genTraversal :: Type -> Type -> IO (Subroutine CompiledTraversal)
+genTraversal :: HasCallStack => Type -> Type -> IO (Subroutine CompiledTraversal)
 genTraversal key_ty val_ty =
   fmap snd $ generate Optimised $ \handler clause_begin key_end clause_end -> do
     traversal handler clause_begin key_end key_ty