Adding parsing of comments to golang (#610)

cpg-library/src/main/golang/frontend/frontend.go

@@ -42,8 +42,9 @@ var env *jnigi.Env
 
 type GoLanguageFrontend struct {
 	*jnigi.ObjectRef
-	File   *ast.File
-	Module *modfile.File
+	File       *ast.File
+	Module     *modfile.File
+	CommentMap ast.CommentMap
 }
 
 func InitEnv(e *jnigi.Env) {

cpg-library/src/main/golang/frontend/handler.go

@@ -131,19 +131,51 @@ func (this *GoLanguageFrontend) HandleFile(fset *token.FileSet, file *ast.File,
 	return
 }
 
-func (this *GoLanguageFrontend) handleDecl(fset *token.FileSet, decl ast.Decl) *cpg.Declaration {
+// handleComments maps comments from ast.Node to a cpg.Node by using ast.CommentMap.
+func (this *GoLanguageFrontend) handleComments(node *cpg.Node, astNode ast.Node) {
+	this.LogDebug("Handling comments for %+v", astNode)
+
+	var comment = ""
+
+	// Lookup ast node in comment map. One cannot use Filter() because this would actually filter all the comments
+	// that are "below" this AST node as well, e.g. in its children. We only want the comments on the node itself.
+	// Therefore we must convert the CommentMap back into an actual map to access the stored entry for the node.
+	comments, ok := (map[ast.Node][]*ast.CommentGroup)(this.CommentMap)[astNode]
+	if !ok {
+		return
+	}
+
+	for _, c := range comments {
+		text := strings.TrimRight(c.Text(), "\n")
+		comment += text
+	}
+
+	if comment != "" {
+		node.SetComment(comment)
+
+		this.LogDebug("Comments: %+v", comment)
+	}
+}
+
+func (this *GoLanguageFrontend) handleDecl(fset *token.FileSet, decl ast.Decl) (d *cpg.Declaration) {
 	this.LogDebug("Handling declaration (%T): %+v", decl, decl)
 
 	switch v := decl.(type) {
 	case *ast.FuncDecl:
-		return (*cpg.Declaration)(this.handleFuncDecl(fset, v))
+		d = (*cpg.Declaration)(this.handleFuncDecl(fset, v))
 	case *ast.GenDecl:
-		return (*cpg.Declaration)(this.handleGenDecl(fset, v))
+		d = (*cpg.Declaration)(this.handleGenDecl(fset, v))
 	default:
 		this.LogError("Not parsing declaration of type %T yet: %+v", v, v)
 		// no match
-		return nil
+		d = nil
 	}
+
+	if d != nil {
+		this.handleComments((*cpg.Node)(d), decl)
+	}
+
+	return
 }
 
 func (this *GoLanguageFrontend) handleFuncDecl(fset *token.FileSet, funcDecl *ast.FuncDecl) *jnigi.ObjectRef {
@@ -262,6 +294,8 @@ func (this *GoLanguageFrontend) handleFuncDecl(fset *token.FileSet, funcDecl *as
 
 		// add parameter to scope
 		this.GetScopeManager().AddDeclaration((*cpg.Declaration)(p))
+
+		this.handleComments((*cpg.Node)(p), param)
 	}
 
 	// parse body
@@ -335,7 +369,6 @@ func (this *GoLanguageFrontend) handleTypeSpec(fset *token.FileSet, typeDecl *as
 	err := this.LogInfo("Type specifier with name %s and type (%T, %+v)", typeDecl.Name.Name, typeDecl.Type, typeDecl.Type)
 	if err != nil {
 		log.Fatal(err)
-
 	}
 
 	switch v := typeDecl.Type.(type) {
@@ -543,63 +576,73 @@ func (this *GoLanguageFrontend) handleIncDecStmt(fset *token.FileSet, incDecStmt
 	return u
 }
 
-func (this *GoLanguageFrontend) handleStmt(fset *token.FileSet, stmt ast.Stmt) *cpg.Statement {
+func (this *GoLanguageFrontend) handleStmt(fset *token.FileSet, stmt ast.Stmt) (s *cpg.Statement) {
 	this.LogDebug("Handling statement (%T): %+v", stmt, stmt)
 
 	switch v := stmt.(type) {
 	case *ast.ExprStmt:
 		// in our cpg, each expression is also a statement,
 		// so we do not need an expression statement wrapper
-		return (*cpg.Statement)(this.handleExpr(fset, v.X))
+		s = (*cpg.Statement)(this.handleExpr(fset, v.X))
 	case *ast.AssignStmt:
-		return (*cpg.Statement)(this.handleAssignStmt(fset, v))
+		s = (*cpg.Statement)(this.handleAssignStmt(fset, v))
 	case *ast.DeclStmt:
-		return (*cpg.Statement)(this.handleDeclStmt(fset, v))
+		s = (*cpg.Statement)(this.handleDeclStmt(fset, v))
 	case *ast.IfStmt:
-		return (*cpg.Statement)(this.handleIfStmt(fset, v))
+		s = (*cpg.Statement)(this.handleIfStmt(fset, v))
 	case *ast.SwitchStmt:
-		return (*cpg.Statement)(this.handleSwitchStmt(fset, v))
+		s = (*cpg.Statement)(this.handleSwitchStmt(fset, v))
 	case *ast.CaseClause:
-		return (*cpg.Statement)(this.handleCaseClause(fset, v))
+		s = (*cpg.Statement)(this.handleCaseClause(fset, v))
 	case *ast.BlockStmt:
-		return (*cpg.Statement)(this.handleBlockStmt(fset, v))
+		s = (*cpg.Statement)(this.handleBlockStmt(fset, v))
 	case *ast.ForStmt:
-		return (*cpg.Statement)(this.handleForStmt(fset, v))
+		s = (*cpg.Statement)(this.handleForStmt(fset, v))
 	case *ast.ReturnStmt:
-		return (*cpg.Statement)(this.handleReturnStmt(fset, v))
+		s = (*cpg.Statement)(this.handleReturnStmt(fset, v))
 	case *ast.IncDecStmt:
-		return (*cpg.Statement)(this.handleIncDecStmt(fset, v))
+		s = (*cpg.Statement)(this.handleIncDecStmt(fset, v))
 	default:
 		this.LogError("Not parsing statement of type %T yet: %+v", v, v)
+		s = nil
 	}
 
-	return nil
+	if s != nil {
+		this.handleComments((*cpg.Node)(s), stmt)
+	}
+
+	return
 }
 
-func (this *GoLanguageFrontend) handleExpr(fset *token.FileSet, expr ast.Expr) *cpg.Expression {
+func (this *GoLanguageFrontend) handleExpr(fset *token.FileSet, expr ast.Expr) (e *cpg.Expression) {
 	this.LogDebug("Handling expression (%T): %+v", expr, expr)
 
 	switch v := expr.(type) {
 	case *ast.CallExpr:
-		return (*cpg.Expression)(this.handleCallExpr(fset, v))
+		e = (*cpg.Expression)(this.handleCallExpr(fset, v))
 	case *ast.IndexExpr:
-		return (*cpg.Expression)(this.handleIndexExpr(fset, v))
+		e = (*cpg.Expression)(this.handleIndexExpr(fset, v))
 	case *ast.BinaryExpr:
-		return (*cpg.Expression)(this.handleBinaryExpr(fset, v))
+		e = (*cpg.Expression)(this.handleBinaryExpr(fset, v))
 	case *ast.UnaryExpr:
-		return (*cpg.Expression)(this.handleUnaryExpr(fset, v))
+		e = (*cpg.Expression)(this.handleUnaryExpr(fset, v))
 	case *ast.SelectorExpr:
-		return (*cpg.Expression)(this.handleSelectorExpr(fset, v))
+		e = (*cpg.Expression)(this.handleSelectorExpr(fset, v))
 	case *ast.BasicLit:
-		return (*cpg.Expression)(this.handleBasicLit(fset, v))
+		e = (*cpg.Expression)(this.handleBasicLit(fset, v))
 	case *ast.Ident:
-		return (*cpg.Expression)(this.handleIdent(fset, v))
+		e = (*cpg.Expression)(this.handleIdent(fset, v))
 	default:
 		this.LogError("Could not parse expression of type %T: %+v", v, v)
+		// TODO: return an error instead?
+		e = nil
 	}
 
-	// TODO: return an error instead?
-	return nil
+	if e != nil {
+		this.handleComments((*cpg.Node)(e), expr)
+	}
+
+	return
 }
 
 func (this *GoLanguageFrontend) handleAssignStmt(fset *token.FileSet, assignStmt *ast.AssignStmt) (expr *cpg.Expression) {

cpg-library/src/main/golang/lib/cpg/main.go

@@ -28,6 +28,7 @@ package main
 import (
 	"cpg"
 	"cpg/frontend"
+	"go/ast"
 	"go/parser"
 	"go/token"
 
@@ -56,6 +57,7 @@ func Java_de_fraunhofer_aisec_cpg_frontends_golang_GoLanguageFrontend_parseInter
 		),
 		nil,
 		nil,
+		ast.CommentMap{},
 	}
 
 	srcObject := jnigi.WrapJObject(uintptr(arg1), "java/lang/String", false)
@@ -81,11 +83,13 @@ func Java_de_fraunhofer_aisec_cpg_frontends_golang_GoLanguageFrontend_parseInter
 	}
 
 	fset := token.NewFileSet()
-	file, err := parser.ParseFile(fset, string(path.([]byte)), string(src.([]byte)), 0)
+	file, err := parser.ParseFile(fset, string(path.([]byte)), string(src.([]byte)), parser.ParseComments)
 	if err != nil {
 		log.Fatal(err)
 	}
 
+	goFrontend.CommentMap = ast.NewCommentMap(fset, file, file.Comments)
+
 	_, err = goFrontend.ParseModule(string(topLevel.([]byte)))
 	if err != nil {
 		goFrontend.LogError("Error occurred while looking for Go modules file: %v", err)

cpg-library/src/main/golang/node.go

@@ -45,6 +45,10 @@ func (n *Node) SetCode(s string) error {
 	return (*jnigi.ObjectRef)(n).SetField(env, "code", NewString(s))
 }
 
+func (n *Node) SetComment(s string) error {
+	return (*jnigi.ObjectRef)(n).SetField(env, "comment", NewString(s))
+}
+
 func (n *Node) SetLocation(location *PhysicalLocation) error {
 	return (*jnigi.ObjectRef)(n).SetField(env, "location", (*jnigi.ObjectRef)(location))
 }

cpg-library/src/test/java/de/fraunhofer/aisec/cpg/frontends/golang/GoLanguageFrontendTest.kt

@@ -29,6 +29,7 @@ import de.fraunhofer.aisec.cpg.BaseTest
 import de.fraunhofer.aisec.cpg.ExperimentalGolang
 import de.fraunhofer.aisec.cpg.TestUtils
 import de.fraunhofer.aisec.cpg.graph.body
+import de.fraunhofer.aisec.cpg.graph.bodyOrNull
 import de.fraunhofer.aisec.cpg.graph.byNameOrNull
 import de.fraunhofer.aisec.cpg.graph.declarations.*
 import de.fraunhofer.aisec.cpg.graph.statements.*
@@ -713,4 +714,53 @@ class GoLanguageFrontendTest : BaseTest() {
         assertNotNull(call)
         assertTrue(call.invokes.contains(newAwesome))
     }
+
+    @Test
+    fun testComments() {
+        val topLevel = Path.of("src", "test", "resources", "golang")
+        val tu =
+            TestUtils.analyzeAndGetFirstTU(
+                listOf(topLevel.resolve("comment.go").toFile()),
+                topLevel,
+                true
+            ) {
+                it.registerLanguage(
+                    GoLanguageFrontend::class.java,
+                    GoLanguageFrontend.GOLANG_EXTENSIONS
+                )
+            }
+
+        assertNotNull(tu)
+
+        val mainNamespace = tu.byNameOrNull<NamespaceDeclaration>("main")
+        assertNotNull(mainNamespace)
+
+        val main = mainNamespace.byNameOrNull<FunctionDeclaration>("main")
+        assertNotNull(main)
+        assertEquals("comment before function", main.comment)
+
+        val i = main.parameters.firstOrNull { it.name == "i" }
+        assertNotNull(i)
+        assertEquals("comment before parameter1", i.comment)
+
+        val j = main.parameters.firstOrNull { it.name == "j" }
+        assertNotNull(j)
+        assertEquals("comment before parameter2", j.comment)
+
+        var declStmt = main.bodyOrNull<DeclarationStatement>()
+        assertNotNull(declStmt)
+        assertEquals("comment before assignment", declStmt.comment)
+
+        declStmt = main.bodyOrNull(1)
+        assertNotNull(declStmt)
+        assertEquals("comment before declaration", declStmt.comment)
+
+        val s = mainNamespace.byNameOrNull<RecordDeclaration>("main.s")
+        assertNotNull(s)
+        assertEquals("comment before struct", s.comment)
+
+        val myField = s.getField("myField")
+        assertNotNull(myField)
+        assertNotNull("comment before field", myField.comment)
+    }
 }

cpg-library/src/test/resources/golang/comment.go

@@ -0,0 +1,19 @@
+package main
+
+// comment before function
+func main(// comment before parameter1
+          i int,
+          // comment before parameter2
+          j int) {
+    // comment before assignment
+    a := "a"
+
+    // comment before declaration
+    var b = a
+}
+
+// comment before struct
+type s struct {
+    // comment before field
+    myField int
+}