From 76cbade86f7478d069b9e3626b3db604cf169251 Mon Sep 17 00:00:00 2001
From: Simon Marlow <smarlow@meta.com>
Date: Wed, 9 Oct 2024 09:50:35 -0700
Subject: [PATCH] Small reordering tweak

Summary:
Classify lookup statements as `StmtPointMatch`. Normally lookups are
tracked separately, but we might encounter lookups inside disjunctions
or nested groups, where classifying them correctly is important.

Reviewed By: malanka

Differential Revision: D64036715

fbshipit-source-id: c748075f545a4df15d896adda7130240ac657889
---
 glean/db/Glean/Query/Codegen/Types.hs |  5 +++++
 glean/db/Glean/Query/Opt.hs           | 12 ------------
 glean/db/Glean/Query/Reorder.hs       | 15 +++------------
 3 files changed, 8 insertions(+), 24 deletions(-)

diff --git a/glean/db/Glean/Query/Codegen/Types.hs b/glean/db/Glean/Query/Codegen/Types.hs
index 49aaea417..b7fa0b55b 100644
--- a/glean/db/Glean/Query/Codegen/Types.hs
+++ b/glean/db/Glean/Query/Codegen/Types.hs
@@ -10,6 +10,7 @@
 module Glean.Query.Codegen.Types
   ( Match(..)
   , matchVar
+  , isVar
   , Var(..)
   , CodegenQuery
   , QueryWithInfo(..)
@@ -187,6 +188,10 @@ type Expr_ var = Term (Match () var)
 
 type Output = Typed (Register 'BinaryOutputPtr)
 
+isVar :: Pat -> Maybe Var
+isVar (Ref v) = matchVar v
+isVar _ = Nothing
+
 data Var = Var
   { varType :: Type
   , varId :: {-# UNPACK #-}!Int
diff --git a/glean/db/Glean/Query/Opt.hs b/glean/db/Glean/Query/Opt.hs
index 875de7e80..86a55b31b 100644
--- a/glean/db/Glean/Query/Opt.hs
+++ b/glean/db/Glean/Query/Opt.hs
@@ -23,13 +23,6 @@ import qualified Glean.Angle.Types as Type
 import Glean.Angle.Types (FieldDef_(..))
 import Glean.Query.Flatten.Types
 import Glean.Query.Codegen.Types
-  ( matchVar
-  , Match(..)
-  , Var(..)
-  , Generator_(..)
-  , Pat
-  , Generator
-  , QueryWithInfo(..))
 import Glean.Query.Vars
 import Glean.RTS.Term
 import Glean.RTS.Types
@@ -501,11 +494,6 @@ extendGen var gen = do
             IntMap.insert (varId var') gen (optGenerators s) }
           return True
 
-isVar :: Pat -> Maybe Var
-isVar (Ref (MatchVar v)) = Just v
-isVar (Ref (MatchBind v)) = Just v
-isVar _ = Nothing
-
 data OptState = OptState
   { optNextVar :: !Int
     -- ^ for making fresh variables
diff --git a/glean/db/Glean/Query/Reorder.hs b/glean/db/Glean/Query/Reorder.hs
index a1be3293f..b2ecfe857 100644
--- a/glean/db/Glean/Query/Reorder.hs
+++ b/glean/db/Glean/Query/Reorder.hs
@@ -29,17 +29,6 @@ import Compat.Prettyprinter hiding ((<>))
 
 import Glean.Query.BindOrder
 import Glean.Query.Codegen.Types
-  ( matchVar
-  , Var(..)
-  , Pat
-  , Match(..)
-  , CgStatement_(..)
-  , CgStatement
-  , Generator_(..)
-  , CgQuery(..)
-  , QueryWithInfo(..)
-  , SeekSection(..)
-  , CodegenQuery)
 import Glean.Display
 import Glean.Query.Flatten.Types
 import Glean.Query.Vars
@@ -319,7 +308,9 @@ reorderStmtGroup sc bound stmts =
       others = [ (bound, stmt) | (Nothing, bound, stmt) <- summaries ]
 
     classify :: VarSet -> FlatStatement -> StmtCost
-    classify bound (FlatStatement _ _ (FactGenerator _ key _ _)) =
+    classify bound (FlatStatement _ lhs (FactGenerator _ key _ _))
+      | Just (Var _ v _) <- isVar lhs, v `IntSet.member` bound = StmtPointMatch
+      | otherwise =
       case classifyPattern ((`IntSet.member` bound) . varId) key of
         PatternMatch _ Point -> StmtPointMatch
         PatternMatch PrefixFactId Scan -> StmtPrefixFactMatch