/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.beam.runners.dataflow.util;

import java.io.Serializable;
import java.util.Map;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.transforms.OldDoFn;
import org.apache.beam.sdk.util.WindowingStrategy;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;

/**
 * Wrapper class holding the necessary information to serialize a {@link OldDoFn}.
 *
 * @param <InputT> the type of the (main) input elements of the {@link OldDoFn}
 * @param <OutputT> the type of the (main) output elements of the {@link OldDoFn}
 */
public class DoFnInfo<InputT, OutputT> implements Serializable {
  private final OldDoFn<InputT, OutputT> doFn;
  private final WindowingStrategy<?, ?> windowingStrategy;
  private final Iterable<PCollectionView<?>> sideInputViews;
  private final Coder<InputT> inputCoder;
  private final long mainOutput;
  private final Map<Long, TupleTag<?>> outputMap;

  public DoFnInfo(OldDoFn<InputT, OutputT> doFn,
      WindowingStrategy<?, ?> windowingStrategy,
      Iterable<PCollectionView<?>> sideInputViews,
      Coder<InputT> inputCoder,
      long mainOutput,
      Map<Long, TupleTag<?>> outputMap) {
    this.doFn = doFn;
    this.windowingStrategy = windowingStrategy;
    this.sideInputViews = sideInputViews;
    this.inputCoder = inputCoder;
    this.mainOutput = mainOutput;
    this.outputMap = outputMap;
  }

  public OldDoFn<InputT, OutputT> getDoFn() {
    return doFn;
  }

  public WindowingStrategy<?, ?> getWindowingStrategy() {
    return windowingStrategy;
  }

  public Iterable<PCollectionView<?>> getSideInputViews() {
    return sideInputViews;
  }

  public Coder<InputT> getInputCoder() {
    return inputCoder;
  }

  public long getMainOutput() {
    return mainOutput;
  }

  public Map<Long, TupleTag<?>> getOutputMap() {
    return outputMap;
  }
}

@@ -847,16 +847,80 @@
 .util;%0A%0A
+import static com.google.common.base.Preconditions.checkState;%0A%0A
 import j
@@ -940,16 +940,16 @@
 izable;%0A
-
 import j
@@ -1003,16 +1003,60 @@
 .Coder;%0A
+import org.apache.beam.sdk.transforms.DoFn;%0A
 import o
@@ -1326,16 +1326,35 @@
 OldDoFn%7D
+%0A * or %7B@link DoFn%7D
 .%0A *%0A * 
@@ -1601,32 +1601,20 @@
 nal 
-OldDoFn%3CInputT, OutputT%3E
+Serializable
  doF
@@ -1868,48 +1868,169 @@
 %0A%0A  
-public DoFnInfo(OldDoFn%3CInputT, OutputT%3E
+/**%0A   * Creates a %7B@link DoFnInfo%7D for the given %7B@link DoFn%7D or %7B@link OldDoFn%7D and auxiliary bits and%0A   * pieces.%0A   */%0A  public DoFnInfo(%0A      Serializable
  doF
@@ -2452,62 +2452,902 @@
 %0A%0A  
-public OldDoFn%3CInputT, OutputT%3E getDoFn() %7B%0A    return
+/**%0A   * @deprecated call the constructor with a %7B@link Serializable%7D%0A   */%0A  @Deprecated%0A  public DoFnInfo(%0A      OldDoFn doFn,%0A      WindowingStrategy%3C?, ?%3E windowingStrategy,%0A      Iterable%3CPCollectionView%3C?%3E%3E sideInputViews,%0A      Coder%3CInputT%3E inputCoder,%0A      long mainOutput,%0A      Map%3CLong, TupleTag%3C?%3E%3E outputMap) %7B%0A    this((Serializable) doFn, windowingStrategy, sideInputViews, inputCoder, mainOutput, outputMap);%0A  %7D%0A%0A  /** Returns the embedded serialized function. It may be a %7B@code DoFn%7D or %7B@code OldDoFn%7D. */%0A  public Serializable getFn() %7B%0A    return doFn;%0A  %7D%0A%0A  /** @deprecated use %7B@link #getFn()%7D */%0A  @Deprecated%0A  public OldDoFn getDoFn() %7B%0A    checkState(%0A        doFn instanceof OldDoFn,%0A        %22Deprecated %25s.getDoFn() called when the payload was actually a new %25s%22,%0A        DoFnInfo.class.getSimpleName(),%0A        DoFn.class.getSimpleName());%0A    return (OldDoFn)
  doF

Make DoFnInfo carry OldDoFn or DoFn