helper/schema: sets properly take into account the diff

helper/schema/resource_data.go

@@ -43,10 +43,11 @@ type getSource byte
 const (
 	getSourceState getSource = 1 << iota
 	getSourceConfig
-	getSourceDiff
 	getSourceSet
-	getSourceExact
-	getSourceMax = getSourceSet
+	getSourceExact               // Only get from the _exact_ level
+	getSourceDiff                // Apply the diff on top our level
+	getSourceLevelMask getSource = getSourceState | getSourceConfig | getSourceSet
+	getSourceMax       getSource = getSourceSet
 )
 
 // getResult is the internal structure that is generated when a Get
@@ -82,7 +83,7 @@ func (d *ResourceData) Get(key string) interface{} {
 // set and the new value is. This is common, for example, for boolean
 // fields which have a zero value of false.
 func (d *ResourceData) GetChange(key string) (interface{}, interface{}) {
-	o, n := d.getChange(key, getSourceConfig, getSourceDiff)
+	o, n := d.getChange(key, getSourceConfig, getSourceConfig|getSourceDiff)
 	return o.Value, n.Value
 }
 
@@ -93,7 +94,7 @@ func (d *ResourceData) GetChange(key string) (interface{}, interface{}) {
 // The first result will not necessarilly be nil if the value doesn't exist.
 // The second result should be checked to determine this information.
 func (d *ResourceData) GetOk(key string) (interface{}, bool) {
-	r := d.getRaw(key, getSourceSet)
+	r := d.getRaw(key, getSourceSet|getSourceDiff)
 	return r.Value, r.Exists
 }
 
@@ -289,12 +290,21 @@ func (d *ResourceData) getSet(
 	// entire set must come from set, diff, state, etc. So we go backwards
 	// and once we get a result, we take it. Or, we never get a result.
 	var raw getResult
-	for listSource := source; listSource > 0; listSource >>= 1 {
-		if source&getSourceExact != 0 && listSource != source {
+	sourceLevel := source & getSourceLevelMask
+	sourceFlags := source & ^getSourceLevelMask
+	for listSource := sourceLevel; listSource > 0; listSource >>= 1 {
+		// If we're already asking for an exact source and it doesn't
+		// match, then leave since the original source was the match.
+		if sourceFlags&getSourceExact != 0 && listSource != sourceLevel {
 			break
 		}
 
-		raw = d.getList(k, nil, schema, listSource|getSourceExact)
+		// The source we get from is the level we're on, plus the flags
+		// we had, plus the exact flag.
+		getSource := listSource
+		getSource |= sourceFlags
+		getSource |= getSourceExact
+		raw = d.getList(k, nil, schema, getSource)
 		if raw.Exists {
 			break
 		}
@@ -384,11 +394,13 @@ func (d *ResourceData) getMap(
 	resultSet := false
 	prefix := k + "."
 
-	exact := source&getSourceExact != 0
-	source &^= getSourceExact
+	flags := source & ^getSourceLevelMask
+	level := source & getSourceLevelMask
+	exact := flags&getSourceExact != 0
+	diff := flags&getSourceDiff != 0
 
-	if !exact || source == getSourceState {
-		if d.state != nil && source >= getSourceState {
+	if !exact || level == getSourceState {
+		if d.state != nil && level >= getSourceState {
 			for k, _ := range d.state.Attributes {
 				if !strings.HasPrefix(k, prefix) {
 					continue
@@ -401,7 +413,7 @@ func (d *ResourceData) getMap(
 		}
 	}
 
-	if d.config != nil && source == getSourceConfig {
+	if d.config != nil && level == getSourceConfig {
 		// For config, we always set the result to exactly what was requested
 		if mraw, ok := d.config.Get(k); ok {
 			result = make(map[string]interface{})
@@ -433,27 +445,25 @@ func (d *ResourceData) getMap(
 		}
 	}
 
-	if !exact || source == getSourceDiff {
-		if d.diff != nil && source >= getSourceDiff {
-			for k, v := range d.diff.Attributes {
-				if !strings.HasPrefix(k, prefix) {
-					continue
-				}
-				resultSet = true
+	if d.diff != nil && diff {
+		for k, v := range d.diff.Attributes {
+			if !strings.HasPrefix(k, prefix) {
+				continue
+			}
+			resultSet = true
 
-				single := k[len(prefix):]
+			single := k[len(prefix):]
 
-				if v.NewRemoved {
-					delete(result, single)
-				} else {
-					result[single] = d.getPrimitive(k, nil, elemSchema, source).Value
-				}
+			if v.NewRemoved {
+				delete(result, single)
+			} else {
+				result[single] = d.getPrimitive(k, nil, elemSchema, source).Value
 			}
 		}
 	}
 
-	if !exact || source == getSourceSet {
-		if d.setMap != nil && source >= getSourceSet {
+	if !exact || level == getSourceSet {
+		if d.setMap != nil && level >= getSourceSet {
 			cleared := false
 			for k, _ := range d.setMap {
 				if !strings.HasPrefix(k, prefix) {
@@ -577,8 +587,10 @@ func (d *ResourceData) getPrimitive(
 	var result string
 	var resultProcessed interface{}
 	var resultComputed, resultSet bool
-	exact := source&getSourceExact != 0
-	source &^= getSourceExact
+	flags := source & ^getSourceLevelMask
+	source = source & getSourceLevelMask
+	exact := flags&getSourceExact != 0
+	diff := flags&getSourceDiff != 0
 
 	if !exact || source == getSourceState {
 		if d.state != nil && source >= getSourceState {
@@ -604,28 +616,26 @@ func (d *ResourceData) getPrimitive(
 		resultComputed = d.config.IsComputed(k)
 	}
 
-	if !exact || source == getSourceDiff {
-		if d.diff != nil && source >= getSourceDiff {
-			attrD, ok := d.diff.Attributes[k]
-			if ok {
-				if !attrD.NewComputed {
-					result = attrD.New
-					if attrD.NewExtra != nil {
-						// If NewExtra != nil, then we have processed data as the New,
-						// so we store that but decode the unprocessed data into result
-						resultProcessed = result
+	if d.diff != nil && diff {
+		attrD, ok := d.diff.Attributes[k]
+		if ok {
+			if !attrD.NewComputed {
+				result = attrD.New
+				if attrD.NewExtra != nil {
+					// If NewExtra != nil, then we have processed data as the New,
+					// so we store that but decode the unprocessed data into result
+					resultProcessed = result
 
-						err := mapstructure.WeakDecode(attrD.NewExtra, &result)
-						if err != nil {
-							panic(err)
-						}
+					err := mapstructure.WeakDecode(attrD.NewExtra, &result)
+					if err != nil {
+						panic(err)
 					}
-
-					resultSet = true
-				} else {
-					result = ""
-					resultSet = false
 				}
+
+				resultSet = true
+			} else {
+				result = ""
+				resultSet = false
 			}
 		}
 	}
@@ -1101,14 +1111,14 @@ func (d *ResourceData) stateSingle(
 func (d *ResourceData) stateSource(prefix string) getSource {
 	// If we're not doing a partial apply, then get the set level
 	if !d.partial {
-		return getSourceSet
+		return getSourceSet | getSourceDiff
 	}
 
 	// Otherwise, only return getSourceSet if its in the partial map.
 	// Otherwise we use state level only.
 	for k, _ := range d.partialMap {
 		if strings.HasPrefix(prefix, k) {
-			return getSourceSet
+			return getSourceSet | getSourceDiff
 		}
 	}
 

helper/schema/resource_data_test.go

@@ -527,6 +527,58 @@ func TestResourceDataGet(t *testing.T) {
 
 			Value: []interface{}{80},
 		},
+
+		{
+			Schema: map[string]*Schema{
+				"data": &Schema{
+					Type:     TypeSet,
+					Optional: true,
+					Elem: &Resource{
+						Schema: map[string]*Schema{
+							"index": &Schema{
+								Type:     TypeInt,
+								Required: true,
+							},
+
+							"value": &Schema{
+								Type:     TypeString,
+								Required: true,
+							},
+						},
+					},
+					Set: func(a interface{}) int {
+						m := a.(map[string]interface{})
+						return m["index"].(int)
+					},
+				},
+			},
+
+			State: &terraform.InstanceState{
+				Attributes: map[string]string{
+					"data.#":       "1",
+					"data.0.index": "10",
+					"data.0.value": "50",
+				},
+			},
+
+			Diff: &terraform.InstanceDiff{
+				Attributes: map[string]*terraform.ResourceAttrDiff{
+					"data.0.value": &terraform.ResourceAttrDiff{
+						Old: "50",
+						New: "80",
+					},
+				},
+			},
+
+			Key: "data",
+
+			Value: []interface{}{
+				map[string]interface{}{
+					"index": 10,
+					"value": "80",
+				},
+			},
+		},
 	}
 
 	for i, tc := range cases {