From b69cafdabf82e0b230a46fe76bcc0e60af13d061 Mon Sep 17 00:00:00 2001
From: Simon Marlow <smarlow@meta.com>
Date: Mon, 28 Oct 2024 06:35:23 -0700
Subject: [PATCH] Fix a query optimisation regression

Summary:
The recent changes exposed a regression with this query

```
0 where
  X1 = codemarkup.hack.HackFileEntityXRefLocations.2
    {file="www/RefClass.php"};
  codemarkup.hack.HackFileEntityXRefLocations.2
    {entity={decl={method=hack.MethodDeclaration.6 {}}}}=X1;
  X0=X1
```

This turns out to be pretty tricky to optimise. I'm adding yet more
flexibility to the final reordering pass that lets it pull out
lookups, which solves the problem in this case.

Reviewed By: iamirzhan

Differential Revision: D65059643

fbshipit-source-id: 879a869ea6d620b28c53454b084917fd0d8bfaa9
---
 glean/db/Glean/Query/Reorder.hs | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/glean/db/Glean/Query/Reorder.hs b/glean/db/Glean/Query/Reorder.hs
index af8b98bca..777b250e2 100644
--- a/glean/db/Glean/Query/Reorder.hs
+++ b/glean/db/Glean/Query/Reorder.hs
@@ -586,6 +586,7 @@ reorderStmts stmts = iterate stmts []
   choose _ [one] = ("only", (one, []))
   choose scope stmts = fromMaybe (error "choose") $
     firstResolved <|>
+    firstLookup <|>
     firstNotUnresolved <|>
     firstNotUnresolvedLenient <|>
     fallback
@@ -593,6 +594,9 @@ reorderStmts stmts = iterate stmts []
       firstResolved = ("resolved",) <$>
         find (isResolvedFilter (ifBoundOnly scope) . fst) stmts'
 
+      firstLookup = ("lookup",) <$>
+        find (isLookup (ifBoundOnly scope) . fst) stmts'
+
       -- try to find a statement that's already resolved without
       -- adding any missing generators, and then try again allowing
       -- unbound predicates to be resolved.
@@ -665,6 +669,11 @@ isReadyFilter scope stmt notFilter = case stmt of
         appearInStmts
   _ -> notFilter
 
+isLookup :: InScope -> FlatStatement -> Bool
+isLookup scope (FlatStatement _ lhs FactGenerator{})
+  | Just (Var _ v _) <- isVar lhs, inScopeBound scope v = True
+isLookup _ _ = False
+
 data InScope = InScope
   { unboundPredicates :: Bool
   , isInScope :: Variable -> Bool