static DeduplicationResult ensureSingleProducer(
      QueryablePipeline pipeline,
      Collection<ExecutableStage> stages,
      Collection<PipelineNode.PTransformNode> unfusedTransforms) {
    RunnerApi.Components.Builder unzippedComponents = pipeline.getComponents().toBuilder();

    Multimap<PipelineNode.PCollectionNode, StageOrTransform> pcollectionProducers =
        getProducers(pipeline, stages, unfusedTransforms);
    Multimap<StageOrTransform, PipelineNode.PCollectionNode> requiresNewOutput =
        HashMultimap.create();
    // Create a synthetic PCollection for each of these nodes. The transforms in the runner
    // portion of the graph that creates them should be replaced in the result components. The
    // ExecutableStage must also be rewritten to have updated outputs and transforms.
    for (Map.Entry<PipelineNode.PCollectionNode, Collection<StageOrTransform>> collectionProducer :
        pcollectionProducers.asMap().entrySet()) {
      if (collectionProducer.getValue().size() > 1) {
        for (StageOrTransform producer : collectionProducer.getValue()) {
          requiresNewOutput.put(producer, collectionProducer.getKey());
        }
      }
    }

    Map<ExecutableStage, ExecutableStage> updatedStages = new LinkedHashMap<>();
    Map<String, PipelineNode.PTransformNode> updatedTransforms = new LinkedHashMap<>();
    Multimap<String, PipelineNode.PCollectionNode> originalToPartial = HashMultimap.create();
    for (Map.Entry<StageOrTransform, Collection<PipelineNode.PCollectionNode>>
        deduplicationTargets : requiresNewOutput.asMap().entrySet()) {
      if (deduplicationTargets.getKey().getStage() != null) {
        StageDeduplication deduplication =
            deduplicatePCollections(
                deduplicationTargets.getKey().getStage(),
                deduplicationTargets.getValue(),
                unzippedComponents::containsPcollections);
        for (Entry<String, PipelineNode.PCollectionNode> originalToPartialReplacement :
            deduplication.getOriginalToPartialPCollections().entrySet()) {
          originalToPartial.put(
              originalToPartialReplacement.getKey(), originalToPartialReplacement.getValue());
          unzippedComponents.putPcollections(
              originalToPartialReplacement.getValue().getId(),
              originalToPartialReplacement.getValue().getPCollection());
        }
        updatedStages.put(
            deduplicationTargets.getKey().getStage(), deduplication.getUpdatedStage());
      } else if (deduplicationTargets.getKey().getTransform() != null) {
        PTransformDeduplication deduplication =
            deduplicatePCollections(
                deduplicationTargets.getKey().getTransform(),
                deduplicationTargets.getValue(),
                unzippedComponents::containsPcollections);
        for (Entry<String, PipelineNode.PCollectionNode> originalToPartialReplacement :
            deduplication.getOriginalToPartialPCollections().entrySet()) {
          originalToPartial.put(
              originalToPartialReplacement.getKey(), originalToPartialReplacement.getValue());
          unzippedComponents.putPcollections(
              originalToPartialReplacement.getValue().getId(),
              originalToPartialReplacement.getValue().getPCollection());
        }
        updatedTransforms.put(
            deduplicationTargets.getKey().getTransform().getId(),
            deduplication.getUpdatedTransform());
      } else {
        throw new IllegalStateException(
            String.format(
                "%s with no %s or %s",
                StageOrTransform.class.getSimpleName(),
                ExecutableStage.class.getSimpleName(),
                PipelineNode.PTransformNode.class.getSimpleName()));
      }
    }

    Set<PipelineNode.PTransformNode> introducedFlattens = new LinkedHashSet<>();
    for (Map.Entry<String, Collection<PipelineNode.PCollectionNode>> partialFlattenTargets :
        originalToPartial.asMap().entrySet()) {
      String flattenId =
          SyntheticComponents.uniqueId("unzipped_flatten", unzippedComponents::containsTransforms);
      PTransform flattenPartialPCollections =
          createFlattenOfPartials(
              flattenId, partialFlattenTargets.getKey(), partialFlattenTargets.getValue());
      unzippedComponents.putTransforms(flattenId, flattenPartialPCollections);
      introducedFlattens.add(PipelineNode.pTransform(flattenId, flattenPartialPCollections));
    }

    Components components = unzippedComponents.build();
    return DeduplicationResult.of(components, introducedFlattens, updatedStages, updatedTransforms);
  }
@Test
  public void duplicateOverStagesAndTransforms() {
    /* When both a stage and a runner-executed transform produce a PCollection, all should be
     * replaced with synthetic flattens.
     * original graph:
     *             --> one -> .out \
     * red -> .out |                -> shared -> .out
     *             --------------> /
     *
     * fused graph:
     *             --> [one -> .out -> shared ->] .out
     * red -> .out |
     *             ------------------> shared --> .out
     *
     * deduplicated graph:
     *             --> [one -> .out -> shared ->] .out:0 \
     * red -> .out |                                      -> shared -> .out
     *             -----------------> shared:0 -> .out:1 /
     */
    PCollection redOut = PCollection.newBuilder().setUniqueName("red.out").build();
    PTransform red =
        PTransform.newBuilder()
            .setSpec(
                FunctionSpec.newBuilder()
                    .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN)
                    .build())
            .putOutputs("out", redOut.getUniqueName())
            .build();
    PCollection oneOut = PCollection.newBuilder().setUniqueName("one.out").build();
    PTransform one =
        PTransform.newBuilder()
            .setSpec(
                FunctionSpec.newBuilder()
                    .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN)
                    .build())
            .putInputs("in", redOut.getUniqueName())
            .putOutputs("out", oneOut.getUniqueName())
            .build();
    PCollection sharedOut = PCollection.newBuilder().setUniqueName("shared.out").build();
    PTransform shared =
        PTransform.newBuilder()
            .setSpec(
                FunctionSpec.newBuilder()
                    .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN)
                    .build())
            .putInputs("one", oneOut.getUniqueName())
            .putInputs("red", redOut.getUniqueName())
            .putOutputs("shared", sharedOut.getUniqueName())
            .build();
    PCollection blueOut = PCollection.newBuilder().setUniqueName("blue.out").build();
    PTransform blue =
        PTransform.newBuilder()
            .setSpec(
                FunctionSpec.newBuilder()
                    .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN)
                    .build())
            .putInputs("in", sharedOut.getUniqueName())
            .putOutputs("out", blueOut.getUniqueName())
            .build();
    RunnerApi.Components components =
        Components.newBuilder()
            .putTransforms("one", one)
            .putPcollections(oneOut.getUniqueName(), oneOut)
            .putTransforms("red", red)
            .putPcollections(redOut.getUniqueName(), redOut)
            .putTransforms("shared", shared)
            .putPcollections(sharedOut.getUniqueName(), sharedOut)
            .putTransforms("blue", blue)
            .putPcollections(blueOut.getUniqueName(), blueOut)
            .build();
    PTransformNode sharedTransform = PipelineNode.pTransform("shared", shared);
    ExecutableStage oneStage =
        ImmutableExecutableStage.of(
            components,
            Environment.getDefaultInstance(),
            PipelineNode.pCollection(redOut.getUniqueName(), redOut),
            ImmutableList.of(),
            ImmutableList.of(),
            ImmutableList.of(),
            ImmutableList.of(PipelineNode.pTransform("one", one), sharedTransform),
            ImmutableList.of(PipelineNode.pCollection(sharedOut.getUniqueName(), sharedOut)),
            DEFAULT_WIRE_CODER_SETTINGS);
    PTransformNode redTransform = PipelineNode.pTransform("red", red);
    PTransformNode blueTransform = PipelineNode.pTransform("blue", blue);
    QueryablePipeline pipeline = QueryablePipeline.forPrimitivesIn(components);
    OutputDeduplicator.DeduplicationResult result =
        OutputDeduplicator.ensureSingleProducer(
            pipeline,
            ImmutableList.of(oneStage),
            ImmutableList.of(redTransform, blueTransform, sharedTransform));

    assertThat(result.getIntroducedTransforms(), hasSize(1));
    PTransformNode introduced = getOnlyElement(result.getIntroducedTransforms());
    assertThat(introduced.getTransform().getOutputsMap().size(), equalTo(1));
    assertThat(
        getOnlyElement(introduced.getTransform().getOutputsMap().values()),
        equalTo(sharedOut.getUniqueName()));

    assertThat(
        result.getDeduplicatedComponents().getPcollectionsMap().keySet(),
        hasItems(introduced.getTransform().getInputsMap().values().toArray(new String[0])));

    assertThat(result.getDeduplicatedStages().keySet(), hasSize(1));
    assertThat(result.getDeduplicatedTransforms().keySet(), containsInAnyOrder("shared"));

    List<String> introducedOutputs = new ArrayList<>();
    introducedOutputs.addAll(
        result.getDeduplicatedTransforms().get("shared").getTransform().getOutputsMap().values());
    introducedOutputs.addAll(
        result.getDeduplicatedStages().get(oneStage).getOutputPCollections().stream()
            .map(PCollectionNode::getId)
            .collect(Collectors.toList()));
    assertThat(
        introduced.getTransform().getInputsMap().values(),
        containsInAnyOrder(introducedOutputs.toArray(new String[0])));
    assertThat(
        result.getDeduplicatedComponents().getPcollectionsMap().keySet(),
        hasItems(introducedOutputs.toArray(new String[0])));
    assertThat(
        result.getDeduplicatedComponents().getTransformsMap(),
        hasEntry(introduced.getId(), introduced.getTransform()));
  }