From 1a79b13bdc93a2be2d51350544e9edeba125af94 Mon Sep 17 00:00:00 2001
From: Erik Krogh Kristensen <erik-krogh@github.com>
Date: Fri, 15 Oct 2021 13:31:35 +0200
Subject: [PATCH] fix performance

---
 ql/src/codeql_ql/ast/Ast.qll             |  5 ++++-
 ql/src/codeql_ql/ast/internal/Module.qll | 16 +++++++++-------
 2 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/ql/src/codeql_ql/ast/Ast.qll b/ql/src/codeql_ql/ast/Ast.qll
index d8cba60276c..81a020b8af1 100644
--- a/ql/src/codeql_ql/ast/Ast.qll
+++ b/ql/src/codeql_ql/ast/Ast.qll
@@ -70,7 +70,10 @@ class AstNode extends TAstNode {
   predicate hasAnnotation(string name) { this.getAnAnnotation().getName() = name }
 
   /** Gets an annotation of this AST node. */
-  Annotation getAnAnnotation() { toQL(this).getParent() = toQL(result).getParent() }
+  cached
+  Annotation getAnAnnotation() {
+    toQL(this).getParent() = pragma[only_bind_out](toQL(result)).getParent()
+  }
 
   /**
    * Gets the predicate that contains this AST node.
diff --git a/ql/src/codeql_ql/ast/internal/Module.qll b/ql/src/codeql_ql/ast/internal/Module.qll
index e5edc0ee875..7ca7abaecbc 100644
--- a/ql/src/codeql_ql/ast/internal/Module.qll
+++ b/ql/src/codeql_ql/ast/internal/Module.qll
@@ -148,14 +148,16 @@ private predicate resolveSelectionName(Import imp, ContainerOrModule m, int i) {
 
 cached
 private module Cached {
-  // TODO: Use `AstNode::getParent` once it is total
-  private QL::AstNode parent(QL::AstNode n) {
-    result = n.getParent() and
-    not n instanceof QL::Module
-  }
-
   private Module getEnclosingModule0(AstNode n) {
-    AstNodes::toQL(result) = parent*(AstNodes::toQL(n).getParent())
+    not n instanceof Module and
+    (
+      n = result.getAChild(_)
+      or
+      exists(AstNode prev |
+        result = getEnclosingModule0(prev) and
+        n = prev.getAChild(_)
+      )
+    )
   }
 
   cached