Make result variables aware of their index.

Author: Max Schaefer

ql/src/semmle/go/Decls.qll

@@ -112,14 +112,7 @@ class FuncDef extends @funcdef, StmtParent, ExprParent {
   DeferStmt getADeferStmt() { result.getEnclosingFunction() = this }
 
   /** Gets the `i`th result variable of this function. */
-  ResultVariable getResultVar(int i) {
-    result =
-      rank[i + 1](ResultVariable res, int j, int k |
-        res.getDeclaration() = getTypeExpr().getResultDecl(j).getNameExpr(k)
-      |
-        res order by j, k
-      )
-  }
+  ResultVariable getResultVar(int i) { result.isResultOf(this, i) }
 
   /** Gets a result variable of this function. */
   ResultVariable getAResultVar() { result.getFunction() = this }

ql/src/semmle/go/Scopes.qll

@@ -274,12 +274,28 @@ class ReceiverVariable extends Parameter {
 
 /** A (named) function result variable. */
 class ResultVariable extends DeclaredVariable {
-  FuncDef fn;
+  FuncDef f;
+  int index;
 
-  ResultVariable() { fn.getTypeExpr().getAResultDecl().getNameExpr(_) = this.getDeclaration() }
+  ResultVariable() {
+    exists(FuncTypeExpr tp | tp = f.getTypeExpr() |
+      this =
+        rank[index + 1](DeclaredVariable parm, int j, int k |
+          parm.getDeclaration() = tp.getResultDecl(j).getNameExpr(k)
+        |
+          parm order by j, k
+        )
+    )
+  }
 
   /** Gets the function to which this result variable belongs. */
-  FuncDef getFunction() { result = fn }
+  FuncDef getFunction() { result = f }
+
+  /** Gets the index of this result among all results of the function. */
+  int getIndex() { result = index }
+
+  /** Holds if this is the `i`th result of function `fd`. */
+  predicate isResultOf(FuncDef fd, int i) { fd = f and i = index }
 }
 
 /**

ql/src/semmle/go/dataflow/internal/DataFlowUtil.qll

@@ -236,9 +236,15 @@ abstract class FunctionNode extends Node {
   abstract ReceiverNode getReceiver();
 
   /**
-   * Gets a value returned by the given function via a return statement or an assignment to a result variable.
+   * Gets a value returned by the given function via a return statement or an assignment to a
+   * result variable.
    */
   abstract ResultNode getAResult();
+
+  /**
+   * Gets the data-flow node corresponding to the `i`th result of this function.
+   */
+  ResultNode getResult(int i) { result = getAResult() and result.getIndex() = i }
 }
 
 /** A representation of a function that is declared in the module scope. */
@@ -533,6 +539,9 @@ class ResultNode extends InstructionNode {
     or
     insn.(IR::ReadResultInstruction).reads(fd.getResultVar(i))
   }
+
+  /** Gets the index of this result among all results of the function. */
+  int getIndex() { result = i }
 }
 
 /**