From 2fbbabda2dcee07263e65f90c4304ab92d727c19 Mon Sep 17 00:00:00 2001
From: Harry Maclean <harry@hmac.dev>
Date: Wed, 4 Aug 2021 16:20:32 +0100
Subject: [PATCH] First draft of a jump-to-definition query

TODO: flesh out this message
---
 ql/src/queries/analysis/Definitions.ql        | 100 ++++++++++++++++++
 .../query-tests/analysis/Definitions.expected |   8 ++
 .../query-tests/analysis/Definitions.qlref    |   1 +
 ql/test/query-tests/analysis/Definitions.rb   |  26 +++++
 4 files changed, 135 insertions(+)
 create mode 100644 ql/src/queries/analysis/Definitions.ql
 create mode 100644 ql/test/query-tests/analysis/Definitions.expected
 create mode 100644 ql/test/query-tests/analysis/Definitions.qlref
 create mode 100644 ql/test/query-tests/analysis/Definitions.rb

diff --git a/ql/src/queries/analysis/Definitions.ql b/ql/src/queries/analysis/Definitions.ql
new file mode 100644
index 00000000000..40e63257a65
--- /dev/null
+++ b/ql/src/queries/analysis/Definitions.ql
@@ -0,0 +1,100 @@
+/**
+ * @name Definitions
+ * @description Jump to definition helper query.
+ * @kind definitions
+ * @id rb/jump-to-definition
+ */
+
+/*
+ * TODO:
+ *    - instance and class variables
+ *    - should `Foo.new` point to `Foo#initialize`?
+ */
+
+import ruby
+import codeql_ruby.ast.internal.Module
+import codeql_ruby.controlflow.ControlFlowGraph
+import codeql_ruby.controlflow.CfgNodes
+import codeql_ruby.dataflow.SSA
+
+from DefLoc loc, Location src, Location target, string kind
+where
+  (
+    exists(ConstantReadAccess read, ConstantWriteAccess write | ConstantDefLoc(read, write) = loc |
+      src = read.getLocation() and
+      target = write.getLocation() and
+      kind = "constant"
+    )
+    or
+    exists(MethodCall call, Method meth | LocalMethodLoc(call, meth) = loc |
+      src = call.getLocation() and
+      target = meth.getLocation() and
+      kind = "method"
+    )
+    or
+    exists(VariableReadAccess read, Ssa::WriteDefinition write |
+      LocalVariableLoc(read, write) = loc
+    |
+      src = read.getLocation() and
+      target = write.getLocation() and
+      kind = "variable"
+    )
+  )
+select src, target, kind
+
+/**
+ * Definition location info for different identifiers.
+ * Each branch holds two values that have a `getLocation()` predicate.
+ * The first is the "source" - some usage of an identifier.
+ * The second is the "target" - the definition of that identifier.
+ */
+newtype DefLoc =
+  /** A constant, module or class. */
+  ConstantDefLoc(ConstantReadAccess read, ConstantWriteAccess write) { write = definitionOf(read) } or
+  /** A call to a method that is defined in the same class as the call. */
+  LocalMethodLoc(MethodCall call, Method meth) {
+    meth = lookupMethod(call.getEnclosingModule().getModule(), call.getMethodName()) and
+    call.getReceiver() instanceof Self
+  } or
+  /** A local variable. */
+  LocalVariableLoc(VariableReadAccess read, Ssa::WriteDefinition write) {
+    read = write.getARead().getExpr() and not read.getLocation() = write.getLocation()
+  }
+
+/**
+ * Gets the fully qualified name for a constant, based on the context in which it is defined.
+ *
+ *  For example, given
+ *  ```ruby
+ *  module Foo
+ *    module Bar
+ *      class Baz
+ *      end
+ *    end
+ *  end
+ *  ```
+ *
+ *  the constant `Bar` has the fully qualified name `Foo::Bar::Baz`.
+ */
+string constantQualifiedName(ConstantWriteAccess w) {
+  not exists(ConstantWriteAccess w2 | w2.getAChild() = w) and result = w.getName()
+  or
+  exists(ConstantWriteAccess w2 |
+    w2.getAChild() = w and result = constantQualifiedName(w2) + "::" + w.getName()
+  )
+}
+
+/**
+ * Gets the constant write that defines the given constant.
+ *  Modules often don't have a unique definition, as they are opened multiple times in different
+ *  files. In these cases we arbitrarily pick the definition with the lexicographically least
+ *  location.
+ */
+ConstantWriteAccess definitionOf(ConstantReadAccess r) {
+  result =
+    max(ConstantWriteAccess w |
+      TResolved(constantQualifiedName(w)) = resolveScopeExpr(r)
+    |
+      w order by w.getLocation().toString()
+    )
+}
diff --git a/ql/test/query-tests/analysis/Definitions.expected b/ql/test/query-tests/analysis/Definitions.expected
new file mode 100644
index 00000000000..301596bb3ad
--- /dev/null
+++ b/ql/test/query-tests/analysis/Definitions.expected
@@ -0,0 +1,8 @@
+| Definitions.rb:4:7:4:9 | Definitions.rb@4:7:4:9 | Definitions.rb:7:5:9:7 | Definitions.rb@7:5:9:7 | method |
+| Definitions.rb:8:7:8:7 | Definitions.rb@8:7:8:7 | Definitions.rb:7:11:7:11 | Definitions.rb@7:11:7:11 | variable |
+| Definitions.rb:12:7:12:7 | Definitions.rb@12:7:12:7 | Definitions.rb:3:5:5:7 | Definitions.rb@3:5:5:7 | method |
+| Definitions.rb:20:7:20:7 | Definitions.rb@20:7:20:7 | Definitions.rb:1:1:15:3 | Definitions.rb@1:1:15:3 | constant |
+| Definitions.rb:20:7:20:10 | Definitions.rb@20:7:20:10 | Definitions.rb:2:3:14:5 | Definitions.rb@2:3:14:5 | constant |
+| Definitions.rb:20:18:20:18 | Definitions.rb@20:18:20:18 | Definitions.rb:19:11:19:11 | Definitions.rb@19:11:19:11 | variable |
+| Definitions.rb:26:1:26:1 | Definitions.rb@26:1:26:1 | Definitions.rb:17:1:24:3 | Definitions.rb@17:1:24:3 | constant |
+| Definitions.rb:26:1:26:4 | Definitions.rb@26:1:26:4 | Definitions.rb:18:3:23:5 | Definitions.rb@18:3:23:5 | constant |
diff --git a/ql/test/query-tests/analysis/Definitions.qlref b/ql/test/query-tests/analysis/Definitions.qlref
new file mode 100644
index 00000000000..a8620aaeec6
--- /dev/null
+++ b/ql/test/query-tests/analysis/Definitions.qlref
@@ -0,0 +1 @@
+queries/analysis/Definitions.ql
diff --git a/ql/test/query-tests/analysis/Definitions.rb b/ql/test/query-tests/analysis/Definitions.rb
new file mode 100644
index 00000000000..350ae469e89
--- /dev/null
+++ b/ql/test/query-tests/analysis/Definitions.rb
@@ -0,0 +1,26 @@
+module A
+  class B
+    def f
+      g 1
+    end
+
+    def g x
+      x
+    end
+
+    def h
+      f
+    end
+  end
+end
+
+module C
+  class D
+    def h y
+      A::B.new.g y
+      UnknownClass.some_method
+    end
+  end
+end
+
+C::D.new