feature: Support for ReturnStmt statement (#180)

* 🍺 supported ReturnStmt statement * 🍺 supported ReturnStmt statement
shivasurya · Nov 6, 2024 · 660aed4 · 660aed4
1 parent db8d91e
commit 660aed4
Show file tree

Hide file tree

Showing 6 changed files with 119 additions and 6 deletions.
diff --git a/sourcecode-parser/graph/construct.go b/sourcecode-parser/graph/construct.go
@@ -54,6 +54,7 @@ type Node struct {
 	ContinueStmt         *model.ContinueStmt
 	YieldStmt            *model.YieldStmt
 	AssertStmt           *model.AssertStmt
+	ReturnStmt           *model.ReturnStmt
 }
 
 type Edge struct {
@@ -182,6 +183,21 @@ func parseJavadocTags(commentContent string) *model.Javadoc {
 func buildGraphFromAST(node *sitter.Node, sourceCode []byte, graph *CodeGraph, currentContext *Node, file string) {
 	isJavaSourceFile := isJavaSourceFile(file)
 	switch node.Type() {
+	case "return_statement":
+		returnNode := javalang.ParseReturnStatement(node, sourceCode)
+		uniqueReturnID := fmt.Sprintf("return_%d_%d_%s", node.StartPoint().Row+1, node.StartPoint().Column+1, file)
+		returnStmtNode := &Node{
+			ID:               GenerateSha256(uniqueReturnID),
+			Type:             "ReturnStmt",
+			LineNumber:       node.StartPoint().Row + 1,
+			Name:             "ReturnStmt",
+			IsExternal:       true,
+			CodeSnippet:      node.Content(sourceCode),
+			File:             file,
+			isJavaSourceFile: isJavaSourceFile,
+			ReturnStmt:       returnNode,
+		}
+		graph.AddNode(returnStmtNode)
 	case "assert_statement":
 		assertNode := javalang.ParseAssertStatement(node, sourceCode)
 		uniqueAssertID := fmt.Sprintf("assert_%d_%d_%s", node.StartPoint().Row+1, node.StartPoint().Column+1, file)

diff --git a/sourcecode-parser/graph/construct_test.go b/sourcecode-parser/graph/construct_test.go
@@ -732,9 +732,9 @@ func TestBuildGraphFromAST(t *testing.T) {
 					}
 				}
 			`,
-			expectedNodes:   4,
+			expectedNodes:   5,
 			expectedEdges:   0,
-			expectedTypes:   []string{"class_declaration", "method_declaration", "binary_expression"},
+			expectedTypes:   []string{"class_declaration", "method_declaration", "binary_expression", "ReturnStmt"},
 			unexpectedTypes: []string{"variable_declaration"},
 		},
 		{
@@ -791,9 +791,9 @@ func TestBuildGraphFromAST(t *testing.T) {
 					}
 				}
 			`,
-			expectedNodes:   73,
+			expectedNodes:   74,
 			expectedEdges:   5,
-			expectedTypes:   []string{"class_declaration", "method_declaration", "binary_expression", "comp_expression", "and_expression", "or_expression", "IfStmt", "ForStmt", "WhileStmt", "DoStmt", "BreakStmt", "ContinueStmt", "YieldStmt"},
+			expectedTypes:   []string{"class_declaration", "method_declaration", "binary_expression", "comp_expression", "and_expression", "or_expression", "IfStmt", "ForStmt", "WhileStmt", "DoStmt", "BreakStmt", "ContinueStmt", "YieldStmt", "ReturnStmt"},
 			unexpectedTypes: []string{""},
 		},
 		{
@@ -811,9 +811,9 @@ func TestBuildGraphFromAST(t *testing.T) {
 					}
 				}
 			`,
-			expectedNodes:   3,
+			expectedNodes:   4,
 			expectedEdges:   0,
-			expectedTypes:   []string{"class_declaration", "method_declaration", "block_comment"},
+			expectedTypes:   []string{"class_declaration", "method_declaration", "block_comment", "ReturnStmt"},
 			unexpectedTypes: []string{"variable_declaration", "binary_expression"},
 		},
 		// add testcase for object creation expression

diff --git a/sourcecode-parser/graph/java/parse_statement.go b/sourcecode-parser/graph/java/parse_statement.go
@@ -42,3 +42,11 @@ func ParseAssertStatement(node *sitter.Node, sourcecode []byte) *model.AssertStm
 	}
 	return assertStmt
 }
+
+func ParseReturnStatement(node *sitter.Node, sourcecode []byte) *model.ReturnStmt {
+	returnStmt := &model.ReturnStmt{}
+	if node.Child(1) != nil {
+		returnStmt.Result = &model.Expr{NodeString: node.Child(1).Content(sourcecode)}
+	}
+	return returnStmt
+}
diff --git a/sourcecode-parser/graph/query.go b/sourcecode-parser/graph/query.go
@@ -151,6 +151,10 @@ func (env *Env) GetAssertStmt() *model.AssertStmt {
 	return env.Node.AssertStmt
 }
 
+func (env *Env) GetReturnStmt() *model.ReturnStmt {
+	return env.Node.ReturnStmt
+}
+
 func QueryEntities(graph *CodeGraph, query parser.Query) (nodes [][]*Node, output [][]interface{}) {
 	result := make([][]*Node, 0)
 
@@ -330,6 +334,7 @@ func generateProxyEnv(node *Node, query parser.Query) map[string]interface{} {
 	continueStmt := "ContinueStmt"
 	yieldStmt := "YieldStmt"
 	assertStmt := "AssertStmt"
+	returnStmt := "ReturnStmt"
 
 	// print query select list
 	for _, entity := range query.SelectList {
@@ -394,6 +399,8 @@ func generateProxyEnv(node *Node, query parser.Query) map[string]interface{} {
 			yieldStmt = entity.Alias
 		case "AssertStmt":
 			assertStmt = entity.Alias
+		case "ReturnStmt":
+			returnStmt = entity.Alias
 		}
 	}
 	env := map[string]interface{}{
@@ -556,6 +563,10 @@ func generateProxyEnv(node *Node, query parser.Query) map[string]interface{} {
 			"toString":      proxyenv.ToString,
 			"getAssertStmt": proxyenv.GetAssertStmt,
 		},
+		returnStmt: map[string]interface{}{
+			"toString":      proxyenv.ToString,
+			"getReturnStmt": proxyenv.GetReturnStmt,
+		},
 	}
 	return env
 }

diff --git a/sourcecode-parser/model/stmt.go b/sourcecode-parser/model/stmt.go
@@ -341,3 +341,37 @@ func (assertStmt *AssertStmt) GetMessage() *Expr {
 func (assertStmt *AssertStmt) GetExpr() *Expr {
 	return assertStmt.Expr
 }
+
+type ReturnStmt struct {
+	Stmt
+	Result *Expr
+}
+
+type IReturnStmt interface {
+	GetAPrimaryQlClass() string
+	GetHalsteadID() int
+	GetPP() string
+	ToString() string
+	GetResult() *Expr
+}
+
+func (returnStmt *ReturnStmt) GetAPrimaryQlClass() string {
+	return "ReturnStmt"
+}
+
+func (returnStmt *ReturnStmt) GetHalsteadID() int {
+	// TODO: Implement Halstead ID calculation for ReturnStmt
+	return 0
+}
+
+func (returnStmt *ReturnStmt) GetPP() string {
+	return fmt.Sprintf("return %s", returnStmt.Result.NodeString)
+}
+
+func (returnStmt *ReturnStmt) ToString() string {
+	return fmt.Sprintf("return %s", returnStmt.Result.NodeString)
+}
+
+func (returnStmt *ReturnStmt) GetResult() *Expr {
+	return returnStmt.Result
+}
diff --git a/sourcecode-parser/model/stmt_test.go b/sourcecode-parser/model/stmt_test.go
@@ -384,3 +384,47 @@ func TestAssertStmt_GetMessage(t *testing.T) {
 		assert.Equal(t, "Modified message", retrievedMessage.NodeString)
 	})
 }
+
+func TestReturnStmt_GetPP(t *testing.T) {
+	t.Run("GetPP with numeric value", func(t *testing.T) {
+		returnStmt := &ReturnStmt{
+			Result: &Expr{NodeString: "42"},
+		}
+		assert.Equal(t, "return 42", returnStmt.GetPP())
+	})
+
+	t.Run("GetPP with string literal", func(t *testing.T) {
+		returnStmt := &ReturnStmt{
+			Result: &Expr{NodeString: "\"hello world\""},
+		}
+		assert.Equal(t, "return \"hello world\"", returnStmt.GetPP())
+	})
+
+	t.Run("GetPP with method call", func(t *testing.T) {
+		returnStmt := &ReturnStmt{
+			Result: &Expr{NodeString: "getValue()"},
+		}
+		assert.Equal(t, "return getValue()", returnStmt.GetPP())
+	})
+
+	t.Run("GetPP with complex expression", func(t *testing.T) {
+		returnStmt := &ReturnStmt{
+			Result: &Expr{NodeString: "x + y * (z - 1)"},
+		}
+		assert.Equal(t, "return x + y * (z - 1)", returnStmt.GetPP())
+	})
+
+	t.Run("GetPP with empty expression", func(t *testing.T) {
+		returnStmt := &ReturnStmt{
+			Result: &Expr{NodeString: ""},
+		}
+		assert.Equal(t, "return ", returnStmt.GetPP())
+	})
+
+	t.Run("GetPP with boolean expression", func(t *testing.T) {
+		returnStmt := &ReturnStmt{
+			Result: &Expr{NodeString: "x > 0 && y < 10"},
+		}
+		assert.Equal(t, "return x > 0 && y < 10", returnStmt.GetPP())
+	})
+}