feature: 🍺 Support for YieldStmt statement (#176)

shivasurya · web-flow · commit 776da67bb446 · 2024-11-04T20:20:34.000-05:00
* 🐳 added YieldStmt support

* 🐛 fix lint

* ⛵ add test coverage
diff --git a/sourcecode-parser/graph/construct.go b/sourcecode-parser/graph/construct.go
@@ -52,7 +52,8 @@ type Node struct {
 	ForStmt              *model.ForStmt
 	BreakStmt            *model.BreakStmt
 	ContinueStmt         *model.ContinueStmt
-} //
+	YieldStmt            *model.YieldStmt
+}
 
 type Edge struct {
 	From *Node
@@ -180,6 +181,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 "yield_statement":
+		yieldNode := javalang.ParseYieldStatement(node, sourceCode)
+		uniqueyieldID := fmt.Sprintf("yield_%d_%d_%s", node.StartPoint().Row+1, node.StartPoint().Column+1, file)
+		yieldStmtNode := &Node{
+			ID:               GenerateSha256(uniqueyieldID),
+			Type:             "YieldStmt",
+			LineNumber:       node.StartPoint().Row + 1,
+			Name:             "YieldStmt",
+			IsExternal:       true,
+			CodeSnippet:      node.Content(sourceCode),
+			File:             file,
+			isJavaSourceFile: isJavaSourceFile,
+			YieldStmt:        yieldNode,
+		}
+		graph.AddNode(yieldStmtNode)
 	case "break_statement":
 		breakNode := javalang.ParseBreakStatement(node, sourceCode)
 		uniquebreakstmtID := fmt.Sprintf("breakstmt_%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
@@ -766,6 +766,19 @@ func TestBuildGraphFromAST(t *testing.T) {
 							System.out.println(i);
                              break;
 						}
+						switch (day) {
+									case "MONDAY" -> 1;
+									case "TUESDAY" -> 2;
+									case "WEDNESDAY" -> 3;
+									case "THURSDAY" -> 4;
+									case "FRIDAY" -> 5;
+									case "SATURDAY" -> 6;
+									case "SUNDAY" -> 7;
+									default -> {
+										System.out.println("Invalid day: " + day);
+										yield 9;  // Using 'yield' to return a value from this case
+									}
+						};
 						do {
 							System.out.println("Hello, World!");
 						} while (a > 0);
@@ -778,9 +791,9 @@ func TestBuildGraphFromAST(t *testing.T) {
 					}
 				}
 			`,
-			expectedNodes:   69,
-			expectedEdges:   4,
-			expectedTypes:   []string{"class_declaration", "method_declaration", "binary_expression", "comp_expression", "and_expression", "or_expression", "IfStmt", "ForStmt", "WhileStmt", "DoStmt", "BreakStmt", "ContinueStmt"},
+			expectedNodes:   73,
+			expectedEdges:   5,
+			expectedTypes:   []string{"class_declaration", "method_declaration", "binary_expression", "comp_expression", "and_expression", "or_expression", "IfStmt", "ForStmt", "WhileStmt", "DoStmt", "BreakStmt", "ContinueStmt", "YieldStmt"},
 			unexpectedTypes: []string{""},
 		},
 		{
diff --git a/sourcecode-parser/graph/java/parse_statement.go b/sourcecode-parser/graph/java/parse_statement.go
@@ -26,3 +26,10 @@ func ParseContinueStatement(node *sitter.Node, sourcecode []byte) *model.Continu
 	}
 	return continueStmt
 }
+
+func ParseYieldStatement(node *sitter.Node, sourcecode []byte) *model.YieldStmt {
+	yieldStmt := &model.YieldStmt{}
+	yieldStmtExpr := &model.Expr{NodeString: node.Child(1).Content(sourcecode)}
+	yieldStmt.Value = yieldStmtExpr
+	return yieldStmt
+}
diff --git a/sourcecode-parser/graph/java/parse_statement_test.go b/sourcecode-parser/graph/java/parse_statement_test.go
@@ -80,3 +80,59 @@ func TestParseContinueStatement(t *testing.T) {
 		})
 	}
 }
+
+func TestParseYieldStatement(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    string
+		expected *model.YieldStmt
+	}{
+		{
+			name:  "Simple yield statement with literal",
+			input: "yield 42;",
+			expected: &model.YieldStmt{
+				Value: &model.Expr{NodeString: "42"},
+			},
+		},
+		{
+			name:  "Yield statement with variable",
+			input: "yield result;",
+			expected: &model.YieldStmt{
+				Value: &model.Expr{NodeString: "result"},
+			},
+		},
+		{
+			name:  "Yield statement with expression",
+			input: "yield a + b;",
+			expected: &model.YieldStmt{
+				Value: &model.Expr{NodeString: "a + b"},
+			},
+		},
+		{
+			name:  "Yield statement with method call",
+			input: "yield getValue();",
+			expected: &model.YieldStmt{
+				Value: &model.Expr{NodeString: "getValue()"},
+			},
+		},
+		{
+			name:  "Yield statement with string literal",
+			input: "yield \"hello\";",
+			expected: &model.YieldStmt{
+				Value: &model.Expr{NodeString: "\"hello\""},
+			},
+		},
+	}
+
+	parser := sitter.NewParser()
+	parser.SetLanguage(java.GetLanguage())
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tree := parser.Parse(nil, []byte(tt.input))
+			node := tree.RootNode().Child(0)
+			result := ParseYieldStatement(node, []byte(tt.input))
+			assert.Equal(t, tt.expected, result)
+		})
+	}
+}
diff --git a/sourcecode-parser/graph/query.go b/sourcecode-parser/graph/query.go
@@ -143,6 +143,10 @@ func (env *Env) GetContinueStmt() *model.ContinueStmt {
 	return env.Node.ContinueStmt
 }
 
+func (env *Env) GetYieldStmt() *model.YieldStmt {
+	return env.Node.YieldStmt
+}
+
 func QueryEntities(graph *CodeGraph, query parser.Query) (nodes [][]*Node, output [][]interface{}) {
 	result := make([][]*Node, 0)
 
@@ -320,6 +324,7 @@ func generateProxyEnv(node *Node, query parser.Query) map[string]interface{} {
 	forStmt := "ForStmt"
 	breakStmt := "BreakStmt"
 	continueStmt := "ContinueStmt"
+	yieldStmt := "YieldStmt"
 
 	// print query select list
 	for _, entity := range query.SelectList {
@@ -380,6 +385,8 @@ func generateProxyEnv(node *Node, query parser.Query) map[string]interface{} {
 			breakStmt = entity.Alias
 		case "ContinueStmt":
 			continueStmt = entity.Alias
+		case "YieldStmt":
+			yieldStmt = entity.Alias
 		}
 	}
 	env := map[string]interface{}{
@@ -534,6 +541,10 @@ func generateProxyEnv(node *Node, query parser.Query) map[string]interface{} {
 			"toString":        proxyenv.ToString,
 			"getContinueStmt": proxyenv.GetContinueStmt,
 		},
+		yieldStmt: map[string]interface{}{
+			"toString":     proxyenv.ToString,
+			"getYieldStmt": proxyenv.GetYieldStmt,
+		},
 	}
 	return env
 }
diff --git a/sourcecode-parser/model/stmt.go b/sourcecode-parser/model/stmt.go
@@ -266,3 +266,38 @@ func (continueStmt *ContinueStmt) hasLabel() bool {
 func (continueStmt *ContinueStmt) GetLabel() string {
 	return continueStmt.Label
 }
+
+// TODO: Implement the SwitchStmt Expr.
+type YieldStmt struct {
+	JumpStmt
+	Value *Expr
+}
+
+type IYieldStmt interface {
+	GetAPrimaryQlClass() string
+	GetHalsteadID() int
+	GetPP() string
+	ToString() string
+	GetValue() *Expr
+}
+
+func (yieldStmt *YieldStmt) GetAPrimaryQlClass() string {
+	return "YieldStmt"
+}
+
+func (yieldStmt *YieldStmt) GetHalsteadID() int {
+	// TODO: Implement Halstead ID calculation for YieldStmt
+	return 0
+}
+
+func (yieldStmt *YieldStmt) GetPP() string {
+	return fmt.Sprintf("yield %s", yieldStmt.Value.NodeString)
+}
+
+func (yieldStmt *YieldStmt) ToString() string {
+	return fmt.Sprintf("yield %s", yieldStmt.Value.NodeString)
+}
+
+func (yieldStmt *YieldStmt) GetValue() *Expr {
+	return yieldStmt.Value
+}
diff --git a/sourcecode-parser/model/stmt_test.go b/sourcecode-parser/model/stmt_test.go
@@ -219,3 +219,133 @@ func TestContinueStmt(t *testing.T) {
 		assert.Equal(t, "", continueStmt.GetLabel())
 	})
 }
+
+func TestYieldStmt(t *testing.T) {
+	t.Run("ToString with non-empty value", func(t *testing.T) {
+		yieldStmt := &YieldStmt{
+			Value: &Expr{NodeString: "42"},
+		}
+		assert.Equal(t, "yield 42", yieldStmt.ToString())
+	})
+
+	t.Run("ToString with empty value", func(t *testing.T) {
+		yieldStmt := &YieldStmt{
+			Value: &Expr{NodeString: ""},
+		}
+		assert.Equal(t, "yield ", yieldStmt.ToString())
+	})
+
+	t.Run("ToString with complex expression", func(t *testing.T) {
+		yieldStmt := &YieldStmt{
+			Value: &Expr{NodeString: "a + b * c"},
+		}
+		assert.Equal(t, "yield a + b * c", yieldStmt.ToString())
+	})
+
+	t.Run("ToString with string literal", func(t *testing.T) {
+		yieldStmt := &YieldStmt{
+			Value: &Expr{NodeString: "\"hello world\""},
+		}
+		assert.Equal(t, "yield \"hello world\"", yieldStmt.ToString())
+	})
+}
+
+func TestYieldStmt_GetValue(t *testing.T) {
+	t.Run("GetValue with non-nil value", func(t *testing.T) {
+		expr := &Expr{NodeString: "42"}
+		yieldStmt := &YieldStmt{
+			Value: expr,
+		}
+		assert.Equal(t, expr, yieldStmt.GetValue())
+	})
+
+	t.Run("GetValue with nil value", func(t *testing.T) {
+		yieldStmt := &YieldStmt{
+			Value: nil,
+		}
+		assert.Nil(t, yieldStmt.GetValue())
+	})
+
+	t.Run("GetValue with complex expression", func(t *testing.T) {
+		expr := &Expr{NodeString: "foo() + bar(x, y)"}
+		yieldStmt := &YieldStmt{
+			Value: expr,
+		}
+		assert.Equal(t, expr, yieldStmt.GetValue())
+	})
+
+	t.Run("GetValue preserves expression reference", func(t *testing.T) {
+		expr := &Expr{NodeString: "someValue"}
+		yieldStmt := &YieldStmt{
+			Value: expr,
+		}
+		retrievedExpr := yieldStmt.GetValue()
+		expr.NodeString = "modifiedValue"
+		assert.Equal(t, "modifiedValue", retrievedExpr.NodeString)
+	})
+}
+
+func TestYieldStmt_GetHalsteadID(t *testing.T) {
+	t.Run("Returns zero for empty yield statement", func(t *testing.T) {
+		yieldStmt := &YieldStmt{}
+		assert.Equal(t, 0, yieldStmt.GetHalsteadID())
+	})
+
+	t.Run("Returns zero for yield with simple value", func(t *testing.T) {
+		yieldStmt := &YieldStmt{
+			Value: &Expr{NodeString: "42"},
+		}
+		assert.Equal(t, 0, yieldStmt.GetHalsteadID())
+	})
+
+	t.Run("Returns zero for yield with complex expression", func(t *testing.T) {
+		yieldStmt := &YieldStmt{
+			Value: &Expr{NodeString: "a + b * c"},
+		}
+		assert.Equal(t, 0, yieldStmt.GetHalsteadID())
+	})
+
+	t.Run("Returns zero for yield with method call", func(t *testing.T) {
+		yieldStmt := &YieldStmt{
+			Value: &Expr{NodeString: "calculateValue()"},
+		}
+		assert.Equal(t, 0, yieldStmt.GetHalsteadID())
+	})
+}
+
+func TestYieldStmt_GetPP(t *testing.T) {
+	t.Run("GetPP with numeric value", func(t *testing.T) {
+		yieldStmt := &YieldStmt{
+			Value: &Expr{NodeString: "42"},
+		}
+		assert.Equal(t, "yield 42", yieldStmt.GetPP())
+	})
+
+	t.Run("GetPP with method call", func(t *testing.T) {
+		yieldStmt := &YieldStmt{
+			Value: &Expr{NodeString: "getValue()"},
+		}
+		assert.Equal(t, "yield getValue()", yieldStmt.GetPP())
+	})
+
+	t.Run("GetPP with complex expression", func(t *testing.T) {
+		yieldStmt := &YieldStmt{
+			Value: &Expr{NodeString: "x + y * (z - 1)"},
+		}
+		assert.Equal(t, "yield x + y * (z - 1)", yieldStmt.GetPP())
+	})
+
+	t.Run("GetPP with empty expression", func(t *testing.T) {
+		yieldStmt := &YieldStmt{
+			Value: &Expr{NodeString: ""},
+		}
+		assert.Equal(t, "yield ", yieldStmt.GetPP())
+	})
+
+	t.Run("GetPP with string literal", func(t *testing.T) {
+		yieldStmt := &YieldStmt{
+			Value: &Expr{NodeString: "\"test string\""},
+		}
+		assert.Equal(t, "yield \"test string\"", yieldStmt.GetPP())
+	})
+}
diff --git a/test-src/android/app/src/main/java/com/ivb/udacity/movieDetailActivity.java b/test-src/android/app/src/main/java/com/ivb/udacity/movieDetailActivity.java
@@ -79,6 +79,18 @@ public boolean onOptionsItemSelected(MenuItem item) {
             i++;
         }
 
+        String message = switch (number) {
+            case ONE -> {
+                yield "Got a 1";
+            }
+            case TWO -> {
+                yield "Got a 2";
+            }
+            default -> {
+                yield a+b;
+            }
+        };
+
         do {
             i++;
         } while (i < 10);

Original file line number	Diff line number	Diff line change
`@@ -26,3 +26,10 @@ func ParseContinueStatement(node sitter.Node, sourcecode []byte) model.Continu`
`26`	`26`	`}`
`27`	`27`	`return continueStmt`
`28`	`28`	`}`
	`29`	`+`
	`30`	`+func ParseYieldStatement(node sitter.Node, sourcecode []byte) model.YieldStmt {`
	`31`	`+ yieldStmt := &model.YieldStmt{}`
	`32`	`+ yieldStmtExpr := &model.Expr{NodeString: node.Child(1).Content(sourcecode)}`
	`33`	`+ yieldStmt.Value = yieldStmtExpr`
	`34`	`+ return yieldStmt`
	`35`	`+}`