From 8ee867758a6925919728950e97bd7717fed0b638 Mon Sep 17 00:00:00 2001
From: sonalmahajan15 <smahajan@uber.com>
Date: Wed, 8 Nov 2023 15:26:20 -0800
Subject: [PATCH] add to unit tests and add assignment flow

---
 annotation/consume_trigger.go      | 25 +++++++++++++++++--------
 annotation/consume_trigger_test.go |  2 ++
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/annotation/consume_trigger.go b/annotation/consume_trigger.go
index 0a5dd1f9..637e159b 100644
--- a/annotation/consume_trigger.go
+++ b/annotation/consume_trigger.go
@@ -886,15 +886,17 @@ func (a *ArgPassDeep) Prestring() Prestring {
 	switch key := a.Ann.(type) {
 	case *ParamAnnotationKey:
 		return ArgPassPrestring{
-			ParamName: key.MinimalString(),
-			FuncName:  key.FuncDecl.Name(),
-			Location:  "",
+			ParamName:     key.MinimalString(),
+			FuncName:      key.FuncDecl.Name(),
+			Location:      "",
+			AssignmentStr: a.assignmentFlow.String(),
 		}
 	case *CallSiteParamAnnotationKey:
 		return ArgPassPrestring{
-			ParamName: key.MinimalString(),
-			FuncName:  key.FuncDecl.Name(),
-			Location:  key.Location.String(),
+			ParamName:     key.MinimalString(),
+			FuncName:      key.FuncDecl.Name(),
+			Location:      key.Location.String(),
+			AssignmentStr: a.assignmentFlow.String(),
 		}
 	default:
 		panic(fmt.Sprintf(
@@ -908,7 +910,8 @@ type ArgPassDeepPrestring struct {
 	FuncName  string
 	// Location points to the code location of the argument pass at the call site for a ArgPass
 	// enclosing CallSiteParamAnnotationKey; Location is empty for a ArgPass enclosing ParamAnnotationKey.
-	Location string
+	Location      string
+	AssignmentStr string
 }
 
 func (a ArgPassDeepPrestring) String() string {
@@ -917,6 +920,7 @@ func (a ArgPassDeepPrestring) String() string {
 	if a.Location != "" {
 		sb.WriteString(fmt.Sprintf(" at %s", a.Location))
 	}
+	sb.WriteString(a.AssignmentStr)
 	return sb.String()
 }
 
@@ -1203,6 +1207,7 @@ func (u *UseAsReturnDeep) Prestring() Prestring {
 		key.RetNum,
 		u.IsNamedReturn,
 		key.FuncDecl.Type().(*types.Signature).Results().At(key.RetNum).Name(),
+		u.assignmentFlow.String(),
 	}
 }
 
@@ -1212,14 +1217,18 @@ type UseAsReturnDeepPrestring struct {
 	RetNum        int
 	IsNamedReturn bool
 	RetName       string
+	AssignmentStr string
 }
 
 func (u UseAsReturnDeepPrestring) String() string {
+	var sb strings.Builder
 	via := ""
 	if u.IsNamedReturn {
 		via = fmt.Sprintf(" via named return `%s`", u.RetName)
 	}
-	return fmt.Sprintf("returned deeply from `%s()`%s in position %d", u.FuncName, via, u.RetNum)
+	sb.WriteString(fmt.Sprintf("returned deeply from `%s()`%s in position %d", u.FuncName, via, u.RetNum))
+	sb.WriteString(u.AssignmentStr)
+	return sb.String()
 }
 
 // overriding position value to point to the raw return statement, which is the source of the potential error
diff --git a/annotation/consume_trigger_test.go b/annotation/consume_trigger_test.go
index b2fc2405..adbf7b98 100644
--- a/annotation/consume_trigger_test.go
+++ b/annotation/consume_trigger_test.go
@@ -58,6 +58,8 @@ var initStructsConsumingAnnotationTrigger = []any{
 	&FldEscape{TriggerIfNonNil: &TriggerIfNonNil{Ann: newMockKey()}},
 	&UseAsNonErrorRetDependentOnErrorRetNilability{TriggerIfNonNil: &TriggerIfNonNil{Ann: newMockKey()}},
 	&UseAsErrorRetWithNilabilityUnknown{TriggerIfNonNil: &TriggerIfNonNil{Ann: newMockKey()}},
+	&ArgPassDeep{TriggerIfDeepNonNil: &TriggerIfDeepNonNil{Ann: newMockKey()}},
+	&UseAsReturnDeep{TriggerIfDeepNonNil: &TriggerIfDeepNonNil{Ann: newMockKey()}},
 }
 
 // ConsumingAnnotationTriggerEqualsTestSuite tests for the `equals` method of all the structs that implement