/*
 * 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.
 */
#ifndef TVM_TIR_SCHEDULE_TRACE_H_
#define TVM_TIR_SCHEDULE_TRACE_H_

#include <tvm/tir/schedule/instruction.h>

namespace tvm {
namespace tir {

// Forward declaration
class Trace;

/*!
 * \brief A callback that allows users to mutate decisions on the fly
 * when applying instructions. The signature of the callback is:
 * \param inst The instruction
 * \param inputs The input random variables
 * \param attrs The attributes
 * \param decision The original decision
 * \return A new decision
 */
using FTraceDecisionProvider = runtime::TypedPackedFunc<ObjectRef(
    const Instruction& inst, const Array<ObjectRef>& inputs, const Array<ObjectRef>& attrs,
    const Optional<ObjectRef>& decision)>;

/*!
 * \brief An execution trace of a scheduling program
 *
 * A trace has two parts:
 * 1) The instructions invoked so far in the program execution
 * 2) The random decisions made upon those instructions, if any
 *
 * A trace can be serialized to:
 * 1) Roundtrippable JSON format: can be saved to file and loaded back
 * 2) Python syntax: allows users to copy-paste the trace to reproduce the scheduling process
 *
 * A trace can be applied to a TensorIR schedule by re-applying all its instructions possibly with
 * their decisions accordingly. Re-sampling is invoked if a sampling instruction doesn't have its
 * corresponding decision; Otherwise the existing decision will be reused accordingly.
 */
class TraceNode : public runtime::Object {
 public:
  /*! \brief The instructions invoked so far in the program execution */
  Array<Instruction> insts;
  /*! \brief The random decisions made upon those instructions */
  Map<Instruction, ObjectRef> decisions;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("insts", &insts);
    v->Visit("decisions", &decisions);
  }

  static constexpr const char* _type_key = "tir.Trace";
  TVM_DECLARE_FINAL_OBJECT_INFO(TraceNode, runtime::Object);

 public:
  /*!
   * \brief Retrieve the decision made on a specific instruction
   * \param inst The instruction whose decision is to be retrieved
   * \return The corresponding decision; NullOpt if there is no decision made on the instruction
   */
  Optional<ObjectRef> GetDecision(const Instruction& inst) const;
  /*!
   * \brief Append a new instruction to the trace
   * \param inst The new instruction to be appended
   */
  void Append(Instruction inst);
  /*!
   * \brief Append a new instruction with a random decision to the trace
   * \param inst The new instruction to be appended
   * \param decision The random decision made on this instruction
   * The type of `decision` depends on the instruction, e.g.
   * the decision of `SamplePerfectTile` has type `Array<IntImm>`
   */
  void Append(Instruction inst, ObjectRef decision);
  /*!
   * \brief Remove the last instruction, along with the decision made on that instruction, if any
   * \return The instruction removed; NullOpt if the trace is empty
   */
  Optional<Instruction> Pop();
  /*!
   * \brief Apply the trace to a TensorIR schedule
   * \param sch The schedule to be applied onto
   * \param remove_postproc If postprocessing instructions are removed
   * \param decision_provider A callback that allows users to mutate decisions on the fly
   * when applying instructions.
   * \sa FTraceDecisionProvider
   */
  void ApplyToSchedule(Schedule sch, bool remove_postproc,
                       FTraceDecisionProvider decision_provider = nullptr) const;
  /*!
   * \brief Serialize the trace as a JSON-style object
   * \param remove_postproc If postprocessing instructions are removed
   * \return The JSON-style object
   */
  ObjectRef AsJSON(bool remove_postproc) const;
  /*!
   * \brief Serialize the trace as a sequence of python statements
   * \param remove_postproc If postprocessing instructions are removed
   * \return A sequence of python statements
   */
  Array<String> AsPython(bool remove_postproc) const;
  /*!
   * \brief Create a new trace with an instruction whose decision is changed,
   * assuming this instruction exists in the resulting trace
   * \param inst The instruction whose decision is to be changed
   * \param decision The decision to be changed to
   * \param remove_postproc If postprocessing instructions are removed
   * \return The new trace with the decision changed
   */
  Trace WithDecision(Instruction inst, ObjectRef decision, bool remove_postproc) const;
  /*!
   * \brief Simplify the trace with dead-code elimination
   * \param remove_postproc If postprocessing instructions are removed
   * \return A simplified trace
   */
  Trace Simplified(bool remove_postproc) const;
};

/*!
 * \brief Managed reference to TraceNode
 * \sa TraceNode
 */
class Trace : public runtime::ObjectRef {
 public:
  /*! \brief Default constructor. Creating an empty trace. */
  Trace();
  /*!
   * \brief Constructor. Creating a trace from existing instructions and their decisions
   * \param insts The instructions used
   * \param decisions The decisions made in sampling
   */
  explicit Trace(Array<Instruction> insts, Map<Instruction, ObjectRef> decisions);
  /*!
   * \brief Apply a JSON-serialized trace to a TensorIR schedule
   * \param json The JSON-serialized trace
   * \param sch The TensorIR schedule
   */
  static void ApplyJSONToSchedule(ObjectRef json, Schedule sch);

  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(Trace, runtime::ObjectRef, TraceNode);
};

}  // namespace tir
}  // namespace tvm

#endif  // TVM_TIR_SCHEDULE_TRACE_H_