removed orphaned data sources from the state

When evaluating orphaned data sources during plan, they need to be
removed from both the refresh state and the current working state.

terraform/context_plan2_test.go

@@ -54,11 +54,6 @@ resource "test_object" "a" {
 			t.Fatalf("expected Create action for missing %s, got %s", c.Addr, c.Action)
 		}
 	}
-
-	_, diags = ctx.Apply()
-	if diags.HasErrors() {
-		t.Fatal(diags.Err())
-	}
 }
 
 func TestContext2Plan_noChangeDataSourceSensitiveNestedSet(t *testing.T) {
@@ -144,3 +139,55 @@ data "test_data_source" "foo" {
 		}
 	}
 }
+
+func TestContext2Plan_orphanDataInstance(t *testing.T) {
+	// ensure the planned replacement of the data source is evaluated properly
+	m := testModuleInline(t, map[string]string{
+		"main.tf": `
+data "test_object" "a" {
+  for_each = { new = "ok" }
+}
+
+output "out" {
+  value = [ for k, _ in data.test_object.a: k ]
+}
+`,
+	})
+
+	p := simpleMockProvider()
+	p.ReadDataSourceFn = func(req providers.ReadDataSourceRequest) (resp providers.ReadDataSourceResponse) {
+		resp.State = req.Config
+		return resp
+	}
+
+	state := states.BuildState(func(s *states.SyncState) {
+		s.SetResourceInstanceCurrent(mustResourceInstanceAddr(`data.test_object.a["old"]`), &states.ResourceInstanceObjectSrc{
+			AttrsJSON: []byte(`{"test_string":"foo"}`),
+			Status:    states.ObjectReady,
+		}, mustProviderConfig(`provider["registry.terraform.io/hashicorp/test"]`))
+	})
+
+	ctx := testContext2(t, &ContextOpts{
+		Config: m,
+		State:  state,
+		Providers: map[addrs.Provider]providers.Factory{
+			addrs.NewDefaultProvider("test"): testProviderFuncFixed(p),
+		},
+	})
+
+	plan, diags := ctx.Plan()
+	if diags.HasErrors() {
+		t.Fatal(diags.Err())
+	}
+
+	change, err := plan.Changes.Outputs[0].Decode()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	expected := cty.TupleVal([]cty.Value{cty.StringVal("new")})
+
+	if change.After.Equals(expected).False() {
+		t.Fatalf("expected %#v, got %#v\n", expected, change.After)
+	}
+}

terraform/node_resource_plan_orphan.go

@@ -60,8 +60,14 @@ func (n *NodePlannableResourceInstanceOrphan) ProvidedBy() (addr addrs.ProviderC
 func (n *NodePlannableResourceInstanceOrphan) dataResourceExecute(ctx EvalContext) tfdiags.Diagnostics {
 	// A data source that is no longer in the config is removed from the state
 	log.Printf("[TRACE] NodePlannableResourceInstanceOrphan: removing state object for %s", n.Addr)
-	state := ctx.RefreshState()
-	state.SetResourceInstanceCurrent(n.Addr, nil, n.ResolvedProvider)
+
+	// we need to update both the refresh state to refresh the current data
+	// source, and the working state for plan-time evaluations.
+	refreshState := ctx.RefreshState()
+	refreshState.SetResourceInstanceCurrent(n.Addr, nil, n.ResolvedProvider)
+
+	workingState := ctx.State()
+	workingState.SetResourceInstanceCurrent(n.Addr, nil, n.ResolvedProvider)
 	return nil
 }