Repository: uber-go/dig (go)

Issue:
Multiple decoration passes within same Scope does not work
When a dependency has multiple passes of decorations being applied within the same Module, it does not work.

For example:
```
c := dig.Container()
c.Provide(func () A { ... })
c.Provide(func (a A) B {...})
c.Provide(func (b B) C { ... })

c.Decorate(func() A { ... })
c.Decorate(func(a A) B { ... })

c.Invoke(func(c C) {
   // should see decorators of A and B getting executed, but will only see decorator of B run.
})
```

Fx version of this would be:
```
fx.Module("module",
  fx.Provide(
    func() A { ... },
    func(a A) B { ... },
    func(b B) C { ... },
  ),
  fx.Decorate(
    func() A { ... },
    func(a A) B { ... },
  ),
)
```

This seems to be a regression from one of the bug fixes that were made to Dig. Will shortly follow up with a fix.

Pull request description:
Fix issue with multiple decoration passes in same Scope
This fixes issue when there is a decorator that depends on a type
that is decorated by another decorator that lives in the same
Scope.

Previously, there is a "circuit-breaker" for the parameter building
logic specifically for the case when the decorator depends on the same
type as its output type; this basically causes the decorator.Call and
paramSingle.Build to call each other, ending up in an infinite recursion.

To avoid that, it was skipping over the current Scope when looking for
decorators for its parameters if it was already in the context of a decorator.

However, that logic is incorrect as it is possible for the decorator can still
depend on *other* types that is decorated within the same Scope.

This fixes the issue by keeping track of the decorator's state in 3 modes -
ready, on stack, and called. If a decorator is on stack, that means that it is
still in the middle of being executed, so when parameter builder finds that
decorator, it knows to skip over the current Scope to look for any decorators
from parent Scopes. Otherwise, it can safely look for decorators from the same
Scope since we won't end up in an infinite recursion.

This fixes #328.

Patch:
diff --git a/decorate.go b/decorate.go
index 2a57b7d..8753440 100644
--- a/decorate.go
+++ b/decorate.go
@@ -28,9 +28,18 @@ import (
 	"go.uber.org/dig/internal/dot"
 )
 
+type decoratorState int
+
+const (
+	decoratorReady decoratorState = iota
+	decoratorOnStack
+	decoratorCalled
+)
+
 type decorator interface {
 	Call(c containerStore) error
 	ID() dot.CtorID
+	State() decoratorState
 }
 
 type decoratorNode struct {
@@ -42,8 +51,8 @@ type decoratorNode struct {
 	// Location where this function was defined.
 	location *digreflect.Func
 
-	// Whether the decorator owned by this node was already called.
-	called bool
+	// Current state of this decorator
+	state decoratorState
 
 	// Parameters of the decorator.
 	params paramList
@@ -87,10 +96,12 @@ func newDecoratorNode(dcor interface{}, s *Scope) (*decoratorNode, error) {
 }
 
 func (n *decoratorNode) Call(s containerStore) error {
-	if n.called {
+	if n.state == decoratorCalled {
 		return nil
 	}
 
+	n.state = decoratorOnStack
+
 	if err := shallowCheckDependencies(s, n.params); err != nil {
 		return errMissingDependencies{
 			Func:   n.location,
@@ -110,12 +121,14 @@ func (n *decoratorNode) Call(s containerStore) error {
 	if err := n.results.ExtractList(n.s, true /* decorated */, results); err != nil {
 		return err
 	}
-	n.called = true
+	n.state = decoratorCalled
 	return nil
 }
 
 func (n *decoratorNode) ID() dot.CtorID { return n.id }
 
+func (n *decoratorNode) State() decoratorState { return n.state }
+
 // DecorateOption modifies the default behavior of Decorate.
 type DecorateOption interface {
 	apply(*decorateOptions)
diff --git a/param.go b/param.go
index 0d7698d..3813e14 100644
--- a/param.go
+++ b/param.go
@@ -227,17 +227,21 @@ func (ps paramSingle) buildWithDecorators(c containerStore, decorating bool) (v
 		decoratingScope containerStore
 	)
 	stores := c.storesToRoot()
-	// If we are already in a decorating stack,
-	// skip the current Scope to avoid infinite
-	// recursion.
-	if decorating {
-		stores = stores[1:]
-	}
+storeLoop:
 	for _, s := range stores {
-		if decorators = s.getValueDecorators(ps.Name, ps.Type); len(decorators) > 0 {
-			decoratingScope = s
-			break
+		if decorators = s.getValueDecorators(ps.Name, ps.Type); len(decorators) == 0 {
+			continue
+		}
+		for _, d := range decorators {
+			if d.State() == decoratorOnStack {
+				// This decorator is already being run.
+				// Avoid a cycle and look further.
+				decorators = nil
+				continue storeLoop
+			}
 		}
+		decoratingScope = s
+		break
 	}
 	if len(decorators) == 0 {
 		return _noValue, false, nil


Test patch:
diff --git a/decorate_test.go b/decorate_test.go
index 7ef0f2c..53febf6 100644
--- a/decorate_test.go
+++ b/decorate_test.go
@@ -124,7 +124,6 @@ func TestDecorateSuccess(t *testing.T) {
 	})
 
 	t.Run("simple decorate a provider to a scope and its descendants", func(t *testing.T) {
-		t.Parallel()
 		type A struct {
 			name string
 		}
@@ -359,6 +358,24 @@ func TestDecorateSuccess(t *testing.T) {
 		}))
 	})
 
+	t.Run("decorate two times within same scope", func(t *testing.T) {
+		type A struct {
+			name string
+		}
+		parent := digtest.New(t)
+		parent.RequireProvide(func() string { return "parent" })
+		parent.RequireProvide(func(n string) A { return A{name: n} })
+
+		child := parent.Scope("child")
+
+		child.RequireDecorate(func() string { return "child" })
+		child.RequireDecorate(func(n string) A { return A{name: n + " decorated"} })
+
+		require.NoError(t, child.Invoke(func(a A) {
+			assert.Equal(t, "child decorated", a.name)
+		}))
+	})
+
 	t.Run("decorate a value group with an empty slice", func(t *testing.T) {
 		type A struct {
 			dig.In