Swift: extract comments

Author: AlexDenisov

swift/codegen/schema.yml

@@ -34,6 +34,10 @@ Location:
   end_column: int
   _pragma: qltest_skip
 
+Comment:
+  _extends: Locatable
+  text: string
+
 Type:
   name: string
   canonical_type: Type

swift/extractor/SwiftExtractor.cpp

@@ -104,11 +104,26 @@ static void extractDeclarations(const SwiftExtractorConfiguration& config,
   trap.emit(unknownFileEntry);
   trap.emit(unknownLocationEntry);
 
+  std::vector<swift::Token> comments;
+  if (primaryFile && primaryFile->getBufferID().hasValue()) {
+    auto& sourceManager = compiler.getSourceMgr();
+    auto tokens = swift::tokenize(compiler.getInvocation().getLangOptions(), sourceManager,
+                                  primaryFile->getBufferID().getValue());
+    for (auto& token : tokens) {
+      if (token.getKind() == swift::tok::comment) {
+        comments.push_back(token);
+      }
+    }
+  }
+
   SwiftVisitor visitor(compiler.getSourceMgr(), trap, module, primaryFile);
   auto topLevelDecls = getTopLevelDecls(module, primaryFile);
   for (auto decl : topLevelDecls) {
     visitor.extract(decl);
   }
+  for (auto& comment : comments) {
+    visitor.extract(comment);
+  }
 }
 
 static std::unordered_set<std::string> collectInputFilenames(swift::CompilerInstance& compiler) {

swift/extractor/infra/SwiftDispatcher.h

@@ -3,6 +3,7 @@
 #include <swift/AST/SourceFile.h>
 #include <swift/Basic/SourceManager.h>
 #include <llvm/Support/FileSystem.h>
+#include <swift/Parse/Token.h>
 
 #include "swift/extractor/trap/TrapLabelStore.h"
 #include "swift/extractor/trap/TrapDomain.h"
@@ -242,6 +243,12 @@ class SwiftDispatcher {
     return false;
   }
 
+  void emitComment(swift::Token& comment) {
+    CommentsTrap entry{trap.createLabel<CommentTag>(), comment.getRawText().str()};
+    trap.emit(entry);
+    attachLocation(comment.getRange().getStart(), comment.getRange().getEnd(), entry.id);
+  }
+
  private:
   void attachLocation(swift::SourceLoc start,
                       swift::SourceLoc end,

swift/extractor/visitors/SwiftVisitor.h

@@ -17,6 +17,7 @@ class SwiftVisitor : private SwiftDispatcher {
   void extract(const T& entity) {
     fetchLabel(entity);
   }
+  void extract(swift::Token& comment) { emitComment(comment); }
 
  private:
   void visit(swift::Decl* decl) override { declVisitor.visit(decl); }

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

@@ -1,6 +1,7 @@
 // generated by codegen/codegen.py
 import codeql.swift.elements.AstNode
 import codeql.swift.elements.Callable
+import codeql.swift.elements.Comment
 import codeql.swift.elements.Element
 import codeql.swift.elements.File
 import codeql.swift.elements.Locatable

swift/ql/lib/codeql/swift/elements/Comment.qll

@@ -0,0 +1,6 @@
+private import codeql.swift.generated.Comment
+
+class Comment extends CommentBase {
+  /** toString */
+  override string toString() { result = getText() }
+}

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

@@ -0,0 +1,8 @@
+// generated by codegen/codegen.py
+import codeql.swift.elements.Locatable
+
+class CommentBase extends @comment, Locatable {
+  override string getAPrimaryQlClass() { result = "Comment" }
+
+  string getText() { comments(this, result) }
+}

swift/ql/lib/swift.dbscheme

@@ -35,6 +35,7 @@ files(
 @locatable =
   @argument
 | @ast_node
+| @comment
 | @condition_element
 | @if_config_clause
 ;
@@ -54,6 +55,11 @@ locations(
   int end_column: int ref
 );
 
+comments(
+  unique int id: @comment,
+  string text: string ref
+);
+
 @type =
   @any_function_type
 | @any_generic_type

swift/ql/test/extractor-tests/comments/comments.expected

@@ -0,0 +1,4 @@
+| comments.swift:1:1:2:1 | // Single line comment\n |
+| comments.swift:3:1:6:3 | /*\n  Multiline\n  comment\n*/ |
+| comments.swift:8:1:9:1 | /// Single line doc comment\n |
+| comments.swift:10:1:13:3 | /**\n  Multiline\n  doc comment\n*/ |

swift/ql/test/extractor-tests/comments/comments.ql

@@ -0,0 +1,4 @@
+import swift
+
+from Comment c
+select c