Swift: extract ObjectLiteralExpr

Author: redsun82

swift/extractor/translators/ExprTranslator.cpp

@@ -526,4 +526,16 @@ codeql::UnresolvedPatternExpr ExprTranslator::translateUnresolvedPatternExpr(
   return entry;
 }
 
+codeql::ObjectLiteralExpr ExprTranslator::translateObjectLiteralExpr(
+    const swift::ObjectLiteralExpr& expr) {
+  auto entry = createExprEntry(expr);
+  entry.kind = static_cast<int>(expr.getLiteralKind());
+  if (auto args = expr.getArgs()) {
+    for (const auto& arg : *args) {
+      entry.arguments.push_back(emitArgument(arg));
+    }
+  }
+  return entry;
+}
+
 }  // namespace codeql

swift/extractor/translators/ExprTranslator.h

@@ -102,6 +102,7 @@ class ExprTranslator : public AstTranslatorBase<ExprTranslator> {
   codeql::ErrorExpr translateErrorExpr(const swift::ErrorExpr& expr);
   codeql::UnresolvedPatternExpr translateUnresolvedPatternExpr(
       const swift::UnresolvedPatternExpr& expr);
+  codeql::ObjectLiteralExpr translateObjectLiteralExpr(const swift::ObjectLiteralExpr& expr);
 
  private:
   void fillAbstractClosureExpr(const swift::AbstractClosureExpr& expr,

swift/ql/lib/codeql/swift/elements/expr/ObjectLiteralExpr.qll

@@ -1,7 +1,21 @@
 private import codeql.swift.generated.expr.ObjectLiteralExpr
 
-class ObjectLiteralExpr extends Generated::ObjectLiteralExpr {
-  override string toString() {
-    result = "#...(...)" // TOOD: We can improve this once we extract the kind
-  }
+class ObjectLiteralExpr extends Generated::ObjectLiteralExpr { }
+
+class FileLiteralExpr extends ObjectLiteralExpr {
+  FileLiteralExpr() { this.getKind() = 0 }
+
+  override string toString() { result = "#fileLiteral(...)" }
+}
+
+class ImageLiteralExpr extends ObjectLiteralExpr {
+  ImageLiteralExpr() { this.getKind() = 1 }
+
+  override string toString() { result = "#imageLiteral(...)" }
+}
+
+class ColorLiteralExpr extends ObjectLiteralExpr {
+  ColorLiteralExpr() { this.getKind() = 2 }
+
+  override string toString() { result = "#colorLiteral(...)" }
 }

swift/ql/lib/codeql/swift/generated/ParentChild.qll

@@ -2588,15 +2588,19 @@ private module Impl {
   private Element getImmediateChildOfObjectLiteralExpr(
     ObjectLiteralExpr e, int index, string partialPredicateCall
   ) {
-    exists(int b, int bLiteralExpr, int n |
+    exists(int b, int bLiteralExpr, int n, int nArgument |
       b = 0 and
       bLiteralExpr =
         b + 1 + max(int i | i = -1 or exists(getImmediateChildOfLiteralExpr(e, i, _)) | i) and
       n = bLiteralExpr and
+      nArgument = n + 1 + max(int i | i = -1 or exists(e.getImmediateArgument(i)) | i) and
       (
         none()
         or
         result = getImmediateChildOfLiteralExpr(e, index - b, partialPredicateCall)
+        or
+        result = e.getImmediateArgument(index - n) and
+        partialPredicateCall = "Argument(" + (index - n).toString() + ")"
       )
     )
   }

swift/ql/lib/codeql/swift/generated/Raw.qll

@@ -839,6 +839,10 @@ module Raw {
 
   class ObjectLiteralExpr extends @object_literal_expr, LiteralExpr {
     override string toString() { result = "ObjectLiteralExpr" }
+
+    int getKind() { object_literal_exprs(this, result) }
+
+    Argument getArgument(int index) { object_literal_expr_arguments(this, index, result) }
   }
 
   class OptionalTryExpr extends @optional_try_expr, AnyTryExpr {

swift/ql/lib/codeql/swift/generated/expr/ObjectLiteralExpr.qll

@@ -1,10 +1,51 @@
 // generated by codegen/codegen.py
 private import codeql.swift.generated.Synth
 private import codeql.swift.generated.Raw
+import codeql.swift.elements.expr.Argument
 import codeql.swift.elements.expr.LiteralExpr
 
 module Generated {
+  /**
+   * `#fileLiteral`, `#imageLiteral` and `#colorLiteral` expressions, which are used in playgrounds.
+   */
   class ObjectLiteralExpr extends Synth::TObjectLiteralExpr, LiteralExpr {
     override string getAPrimaryQlClass() { result = "ObjectLiteralExpr" }
+
+    /**
+     * Gets the
+     *    This is 0 for `#fileLiteral`, 1 for `#imageLiteral` and 2 for `#colorLiteral`.
+     *    .
+     */
+    int getKind() {
+      result = Synth::convertObjectLiteralExprToRaw(this).(Raw::ObjectLiteralExpr).getKind()
+    }
+
+    /**
+     * Gets the `index`th argument of this object literal expression (0-based).
+     *
+     * This includes nodes from the "hidden" AST. It can be overridden in subclasses to change the
+     * behavior of both the `Immediate` and non-`Immediate` versions.
+     */
+    Argument getImmediateArgument(int index) {
+      result =
+        Synth::convertArgumentFromRaw(Synth::convertObjectLiteralExprToRaw(this)
+              .(Raw::ObjectLiteralExpr)
+              .getArgument(index))
+    }
+
+    /**
+     * Gets the `index`th argument of this object literal expression (0-based).
+     */
+    final Argument getArgument(int index) { result = getImmediateArgument(index).resolve() }
+
+    /**
+     * Gets any of the arguments of this object literal expression.
+     */
+    final Argument getAnArgument() { result = getArgument(_) }
+
+    /**
+     * Gets the number of arguments of this object literal expression.
+     */
+    final int getNumberOfArguments() { result = count(getAnArgument()) }
   }
 }

swift/ql/lib/swift.dbscheme

@@ -1307,7 +1307,15 @@ nil_literal_exprs(  //dir=expr
 );
 
 object_literal_exprs(  //dir=expr
-  unique int id: @object_literal_expr
+  unique int id: @object_literal_expr,
+  int kind: int ref
+);
+
+#keyset[id, index]
+object_literal_expr_arguments(  //dir=expr
+  int id: @object_literal_expr ref,
+  int index: int ref,
+  int argument: @argument_or_none ref
 );
 
 optional_try_exprs(  //dir=expr

swift/ql/test/extractor-tests/generated/expr/ObjectLiteralExpr/MISSING_SOURCE.txt

@@ -0,0 +1,3 @@
+| object_literals.swift:3:5:3:42 | #fileLiteral(...) | getKind: | 0 |
+| object_literals.swift:4:5:4:61 | #colorLiteral(...) | getKind: | 2 |
+| object_literals.swift:5:5:5:44 | #imageLiteral(...) | getKind: | 1 |

swift/ql/test/extractor-tests/generated/expr/ObjectLiteralExpr/ObjectLiteralExpr.expected

@@ -0,0 +1,10 @@
+// generated by codegen/codegen.py
+import codeql.swift.elements
+import TestUtils
+
+from ObjectLiteralExpr x, int getKind
+where
+  toBeTested(x) and
+  not x.isUnknown() and
+  getKind = x.getKind()
+select x, "getKind:", getKind