Rust: add location definitions for format arguments

Author: aibaars

rust/ql/lib/ide-contextual-queries/Definitions.qll

@@ -5,13 +5,113 @@
 
 private import codeql.rust.elements.Variable
 private import codeql.rust.elements.Locatable
+private import codeql.rust.elements.FormatArgsExpr
+private import codeql.rust.elements.FormatArgsArg
+private import codeql.rust.elements.Format
+private import codeql.rust.elements.NamedFormatArgument
+private import codeql.rust.elements.PositionalFormatArgument
+private import codeql.Locations
 
-private predicate localVariable(Locatable e, Variable def) { e = def.getAnAccess() }
+/** An element with an associated definition. */
+abstract class Use extends Locatable {
+  /** Gets the definition associated with this element. */
+  abstract Definition getDefinition();
+
+  /**
+   * Gets the type of use.
+   */
+  abstract string getUseType();
+}
+
+private newtype TDef =
+  TVariable(Variable v) or
+  TFormatArgsArgName(Name name) { name = any(FormatArgsArg a).getName() } or
+  TFormatArgsArgIndex(Expr e) { e = any(FormatArgsArg a).getExpr() }
+
+/** A definition */
+class Definition extends TDef {
+  /** Gets the location of this variable. */
+  Location getLocation() {
+    result = this.asVariable().getLocation() or
+    result = this.asName().getLocation() or
+    result = this.asExpr().getLocation()
+  }
+
+  /** Gets this definition as a `Variable` */
+  Variable asVariable() { this = TVariable(result) }
+
+  /** Gets this definition as a `Name` */
+  Name asName() { this = TFormatArgsArgName(result) }
+
+  /** Gets this definition as an `Expr` */
+  Expr asExpr() { this = TFormatArgsArgIndex(result) }
+
+  /** Gets the string representation of this element. */
+  string toString() {
+    result = this.asExpr().toString() or
+    result = this.asVariable().toString() or
+    result = this.asName().getText()
+  }
+}
+
+private class LocalVariableUse extends Use instanceof VariableAccess {
+  private Variable def;
+
+  LocalVariableUse() { this = def.getAnAccess() }
+
+  override Definition getDefinition() { result.asVariable() = def }
+
+  override string getUseType() { result = "local variable" }
+}
+
+private class NamedFormatArgumentUse extends Use instanceof NamedFormatArgument {
+  private Name def;
+
+  NamedFormatArgumentUse() {
+    exists(FormatArgsExpr parent |
+      parent = this.getParent().getParent() and
+      parent.getAnArg().getName() = def and
+      this.getName() = def.getText()
+    )
+  }
+
+  override Definition getDefinition() { result.asName() = def }
+
+  override string getUseType() { result = "format argument" }
+}
+
+private class PositionalFormatUse extends Use instanceof Format {
+  PositionalFormatUse() { not exists(this.getArgumentRef()) }
+
+  override Definition getDefinition() {
+    exists(FormatArgsExpr parent, int index | parent.getFormat(_) = this |
+      this = rank[index + 1](PositionalFormatUse f, int i | parent.getFormat(i) = f | f order by i) and
+      result.asExpr() = parent.getArg(index).getExpr()
+    )
+  }
+
+  override string getUseType() { result = "format argument" }
+}
+
+private class PositionalFormatArgumentUse extends Use instanceof PositionalFormatArgument {
+  private Expr def;
+
+  PositionalFormatArgumentUse() {
+    exists(FormatArgsExpr parent |
+      parent = this.getParent().getParent() and
+      def = parent.getArg(this.getIndex()).getExpr()
+    )
+  }
+
+  override Definition getDefinition() { result.asExpr() = def }
+
+  override string getUseType() { result = "format argument" }
+}
 
 /**
  * Gets an element, of kind `kind`, that element `use` uses, if any.
  */
 cached
-Variable definitionOf(Locatable use, string kind) {
-  localVariable(use, result) and kind = "local variable"
+Definition definitionOf(Use use, string kind) {
+  result = use.getDefinition() and kind = use.getUseType()
 }

rust/ql/lib/ide-contextual-queries/localDefinitions.ql

@@ -8,14 +8,12 @@
  */
 
 import codeql.IDEContextual
-import codeql.rust.elements.Variable
-import codeql.rust.elements.Locatable
 import Definitions
 
 external string selectedSourceFile();
 
-from Locatable e, Variable def, string kind
+from Use use, Definition def, string kind
 where
-  def = definitionOf(e, kind) and
-  e.getLocation().getFile() = getFileBySourceArchiveName(selectedSourceFile())
-select e, def, kind
+  def = definitionOf(use, kind) and
+  use.getLocation().getFile() = getFileBySourceArchiveName(selectedSourceFile())
+select use, def, kind

rust/ql/lib/ide-contextual-queries/localReferences.ql

@@ -8,14 +8,12 @@
  */
 
 import codeql.IDEContextual
-import codeql.rust.elements.Locatable
-import codeql.rust.elements.Variable
 import Definitions
 
 external string selectedSourceFile();
 
-from Locatable e, Variable def, string kind
+from Use use, Definition def, string kind
 where
-  def = definitionOf(e, kind) and
+  def = definitionOf(use, kind) and
   def.getLocation().getFile() = getFileBySourceArchiveName(selectedSourceFile())
-select e, def, kind
+select use, def, kind