text: /*
   * SPDX-License-Identifier: Apache-2.0
   */
  
  #include "ConditionalHelpers.hpp"
  #include "ModelImporter.hpp"
  #include "onnx2trt_utils.hpp"
  #include "toposort.hpp"
  
  namespace onnx2trt
  {
  
  using NodeName = std::string;
  using LayerName = std::string;
  using InputIndex = int32_t;
  
  // A SubgraphPortsMap maps either the inputs or outputs ports of each node in an ONNX graph.
  using SubgraphPortsMap = std::unordered_map<NodeName, std::set<InputIndex>>;
  
  // An InputsMap tracks which IIfConditionalInputLayer we've added to a layer's inputs,
  // so that we can reuse them if needed.
  using InputsMap = std::unordered_map<LayerName, nvinfer1::IIfConditionalInputLayer*>;
  
  // Search for a network Layer name in a SubgraphPortsMap using partial (prefix) name matching.
  // ONNX nodes are matched to network layers using prefix-matching because an ONNX node may have
  // several network layers associcated with it.
  SubgraphPortsMap::const_iterator findLayer(const SubgraphPortsMap& inputs, const std::string layerName)
  {
      return std::find_if(inputs.begin(), inputs.end(), [&](const auto& item) {
          const auto& key = item.first;
          return layerName.compare(0, key.size(), key) == 0;
      });
  }
  
  // Add an ConditionalInputLayer between `layer` and its inputs.
  // I.e. input[inIdx] -> layer ==> input[inIdx] -> ConditionalInputLayer -> layer.
  Status addConditionalInputLayer(IImporterContext* ctx, nvinfer1::IIfConditional* conditional, InputsMap& inputsMap,
      nvinfer1::ILayer& layer, int32_t inIdx)
  {
      auto input = layer.getInput(inIdx);
      if (input == nullptr)
      {
          // Phantom input (an input that is really constant weights).
          return Status::success();
      }
  
      if (layer.getType() == nvinfer1::LayerType::kCONDITIONAL_OUTPUT)
      {
          return Status::success();
      }
  
      auto const name = input->getName();
      auto it = inputsMap.find(name);
      nvinfer1::IIfConditionalInputLayer* inputLayer = nullptr;
      if (it == inputsMap.end())
      {
          inputLayer = conditional->addInput(*input);
          inputsMap[name] = inputLayer;
          const std::string inputLayerName(name);
          ctx->registerLayer(inputLayer, inputLayerName + "_InputLayer");
          ctx->registerTensor(TensorOrWeights{inputLayer->getOutput(0)}, inputLayerName + "_InputLayer_output");
      }
      else
      {
          // An InputLayer may in the inputsMap if it has several consumers.
          inputLayer = it->second;
      }
      layer.setInput(inIdx, *(inputLayer->getOutput(0)));
      return Status::success();
  };
  
  // Take a snapshot of the network before and after parsing the subgraph and return a list
  // of newly added network layers.
  Status importSubgraph(
      IImporterContext* ctx, const ::ONNX_NAMESPACE::GraphProto& subgraph, std::vector<nvinfer1::ILayer*>& newLayers)
  {
      auto net = ctx->network();
      int32_t beforeSubgraph = net->getNbLayers();
      CHECK(onnx2trt::parseGraph(ctx, subgraph));
  
      for (int32_t i = beforeSubgraph; i < net->getNbLayers(); i++)
      {
          newLayers.push_back(net->getLayer(i));
      }
  
      return Status::success();
  }
  
  // Add an IConditionalInputLayer to `layer`'s inputs, if they don't already exist.
  Status addConditionalInputIfNeeded(IImporterContext* ctx, nvinfer1::IIfConditional* conditional, InputsMap& inputsMap,
      nvinfer1::ILayer& layer, SubgraphPortsMap subgraphInputsMap)
  {
      // Return all of the layer's inputs that are external to the subgraph that
      // that the layer belongs to.
      auto getLayerExternalInputs = [&](std::string const& layerName) {
          std::set<int32_t> inIndices;
          auto iter = findLayer(subgraphInputsMap, layerName);
          if (iter != subgraphInputsMap.end())
          {
              const auto& indicesSet = iter->second;
              inIndices.insert(indicesSet.begin(), indicesSet.end());
          }
  
          return inIndices;
      };
  
      const auto inIndices = getLayerExternalInputs(layer.getName());
      for (auto inIdx : inIndices)
      {
          LOG_VERBOSE("Adding Input layer for " << layer.getName());
          addConditionalInputLayer(ctx, conditional, inputsMap, layer, inIdx);
      }
      return Status::success();
  }
  
  // Add IConditionalInputLayers to `layer`'s inputs.
  Status addIfInputLayers(IImporterContext* ctx, nvinfer1::IIfConditional* conditional, InputsMap& inputsMap,
      const std::vector<nvinfer1::ILayer*>& newLayers)
  {
      // Find all of the tensors entering the subgraph.
      // The node-names are from the ONNX context.
      using NodeName = std::string;
      using InputIndex = int32_t;
      std::unordered_map<NodeName, std::set<InputIndex>> subgraphInputsMap;
      getSubgraphInputs(newLayers, subgraphInputsMap);
  
      // Add a ConditionalInputLayer in front of each input that is external to the subgraph.
      for (const auto& layer : newLayers)
      {
          addConditionalInputIfNeeded(ctx, conditional, inputsMap, *layer, subgraphInputsMap);
      }
  
      return Status::success();
  }
  
  // Add an IConditionalOutputLayer to `layer`'s outputs.
  Status addIfOutputLayers(IImporterContext* ctx, nvinfer1::IIfConditional* conditional,
      const ::ONNX_NAMESPACE::GraphProto& thenGraph, const std::vector<nvinfer1::ILayer*>& thenLayers,
      const ::ONNX_NAMESPACE::GraphProto& elseGraph, const std::vector<nvinfer1::ILayer*>& elseLayers,
      std::vector<TensorOrWeights>& graphOutputs)
  {
      // Reported outputs are outputs that the ONNX model reports as subgraph outputs.  This list is
      // not sufficient because it may produce names that are not fully compatible with TensorRT's naming.
      // We use this list to help find the subgraph (SG) output tensors.
      auto getReportedOutputs
          = [&ctx](const ::ONNX_NAMESPACE::GraphProto& body, std::vector<std::string>& reportedOutputs) {
                // Assuming that the subgraph was imported already, we can iterate on its output tensors.
                const auto nbOutputs = body.output_size();
                for (auto i = 0; i < nbOutputs; i++)
                {
                    reportedOutputs.emplace_back(body.output(i).name());
                }
            };
  
      using NodeName = std::string;
      std::unordered_map<NodeName, std::set<int32_t>> thenOutputs;
      std::unordered_map<NodeName, std::set<int32_t>> elseOutputs;
  
      std::vector<std::string> thenReportedOutputs;
      getReportedOutputs(thenGraph, thenReportedOutputs);
      getSubgraphOutputs(thenLayers, thenOutputs, thenReportedOutputs);
      std::vector<std::string> elseReportedOutputs;
      getReportedOutputs(elseGraph, elseReportedOutputs);
      getSubgraphOutputs(elseLayers, elseOutputs, elseReportedOutputs);
  
      // Retrieve the output tensors of a subgraph (tensors exiting the subgraph).
      auto getSubgraphOutputTensors
          = [](IImporterContext* ctx, std::vector<nvinfer1::ITensor*>& sgOutputs, SubgraphPortsMap& subgraphOutputs,
                const ::ONNX_NAMESPACE::GraphProto& subgraph, std::vector<nvinfer1::ILayer*> subgraphLayers) {
                for (const auto& layer : subgraphLayers)
                {
                    const auto layerName = layer->getName();
                    auto iter = findLayer(subgraphOutputs, layerName);
                    if (iter != subgraphOutputs.end())
                    {
                        sgOutputs.push_back(layer->getOutput(0));
                    }
                }
  
                if (sgOutputs.empty())
                {
                    // No new layers, so we can't deduce the outputs and have to use what ONNX tells us.
                    const int32_t nbOutputs = subgraph.output_size();
                    for (int32_t outIdx = 0; outIdx < nbOutputs; outIdx++)
                    {
                        const auto thenName = subgraph.output(outIdx).name();
                        auto* thenTensor = &convertToTensor(ctx->tensors().at(thenName), ctx);
                        sgOutputs.push_back(thenTensor);
                    }
                }
            };
  
      std::vector<nvinfer1::ITensor*> thenOutputTensors;
      getSubgraphOutputTensors(ctx, thenOutputTensors, thenOutputs, thenGraph, thenLayers);
  
      std::vector<nvinfer1::ITensor*> elseSGOutputTensors;
      getSubgraphOutputTensors(ctx, elseSGOutputTensors, elseOutputs, elseGraph, elseLayers);
  
      ASSERT(thenOutputTensors.size() == elseSGOutputTensors.size()
              && "The then/else branches of an If operator must have the same number of outputs.",
          ErrorCode::kINVALID_NODE);
  
      // Add an ConditionalOutputLayer with one output and two inputs
      // (one from the thenGraph and another from the elseGraph).
      for (size_t i = 0; i < elseSGOutputTensors.size(); i++)
      {
          auto* outputLayer = conditional->addOutput(*thenOutputTensors[i], *elseSGOutputTensors[i]);
          ctx->registerLayer(outputLayer, std::string(conditional->getName()) + "_OutputLayer");
          graphOutputs.emplace_back(outputLayer->getOutput(0));
      }
      return Status::success();
  }
  
  // Given a subgraph, find all of its external inputs/outputs (tensors entering/exiting the subgraph).
  Status getSubgraphTensors(const std::vector<nvinfer1::ILayer*>& newLayers,
      std::unordered_map<std::string, std::set<int32_t>>& externalOutputs, bool extractOutputs,
      const std::vector<std::string>* reportedOutputs = nullptr)
  {
      using NodeName = std::string;
      using TensorName = std::string;
      using PortIndex = int32_t;
      using Port = std::pair<NodeName, PortIndex>;
      using TensorsSet = std::unordered_set<nvinfer1::ITensor*>;
      TensorsSet outputTensors;
      TensorsSet inputTensors;
  
      // To determine which tensors are entering or exiting the given graph, we first collect the sets of all input and
      // output tensors. Then we categorize the tensors according to this logic:
      //  Entering tensors := {inputs} - {outputs}
      //  Exiting tensors := {outputs} - {inputs}
  
      // Collect all input and output tensors belonging to nodes in the graph.
  
      auto getTensors = [](nvinfer1::ILayer const* l, bool const input, auto inserter) {
          auto const count = input ? l->getNbInputs() : l->getNbOutputs();
          for (int32_t i = 0; i < count; i++)
          {
              inserter(input ? l->getInput(i) : l->getOutput(i));
          }
      };
  
      for (const auto& l : newLayers)
      {
          getTensors(l, false, [&](nvinfer1::ITensor* t) { outputTensors.insert(t); });
          getTensors(l, true, [&](nvinfer1::ITensor* t) { inputTensors.insert(t); });
      }
  
      using TensorsVec = std::vector<nvinfer1::ITensor*>;
      auto getOutputs = [&](nvinfer1::ILayer const* l, TensorsVec res) {
          getTensors(l, false, [&](nvinfer1::ITensor* t) { res.emplace_back(t); });
      };
  
      auto getInputs = [&](nvinfer1::ILayer const* l, TensorsVec res) {
          getTensors(l, true, [&](nvinfer1::ITensor* t) { res.emplace_back(t); });
      };
  
      // Retrieve the list of tensors either exiting or entering the subgraph.
      std::unordered_map<TensorName, std::vector<Port>> externalPortsMap;
      auto filterTensors = [&](TensorsSet const& tensors, auto getNodeAccessor) {
          for (nvinfer1::ILayer const* l : newLayers)
          {
              const auto& nodeName = l->getName();
              PortIndex i = 0;
  
              TensorsVec nodeAccessor;
              getNodeAccessor(l, nodeAccessor);
              for (const auto& tensor : nodeAccessor)
              {
                  if (tensor == nullptr)
                  {
                      continue;
                  }
                  if (tensors.count(tensor) == 0)
                  {
                      TensorName tensorName = tensor->getName();
                      auto prefixFound = false;
                      if (reportedOutputs)
                      {
                          // reportedOutputs are the names of the outputs as reported by the ONNX parser and help
                          // us further filter the output tensors.
                          //      Exiting tensors := {outputs} - {inputs} - {unreported tensors}
                          // An example: a Split node is internal to a subgraph and has 4 outputs, but only two are
                          // connected to the rest of the graph.  To prevent mistaking the 2 unused outputs as subgraph
                          // outputs, we look for them in reportedOutputs which leads us to ignore the 2 tensors.
                          const auto iter = std::find_if(
                              reportedOutputs->begin(), reportedOutputs->end(), [&](const auto& outputName) {
                                  // Prefix name matching.
                                  return tensorName.compare(0, outputName.size(), outputName) == 0;
                              });
                          prefixFound = iter != reportedOutputs->end();
                      }
                      if (!reportedOutputs || prefixFound)
                      {
                          externalPortsMap[tensorName].push_back(std::make_pair(nodeName, i));
                      }
                  }
                  i++;
              }
          }
      };
  
      if (extractOutputs)
      {
          filterTensors(inputTensors, getOutputs);
      }
      else
      {
          filterTensors(outputTensors, getInputs);
      }
  
      // Create the user's view of the external inputs, which uses the node-name as the key for
      // looking up input/output port index.
      for (auto const& input : externalPortsMap)
      {
          for (const Port& inPort : input.second)
          {
              auto const nodeName = inPort.first;
              auto const portIndex = inPort.second;
              externalOutputs[nodeName].insert(portIndex);
          }
      }
      return Status::success();
  }
  
  Status getSubgraphOutputs(const std::vector<nvinfer1::ILayer*>& newLayers,
      std::unordered_map<std::string, std::set<int32_t>>& externalOutputs,
      const std::vector<std::string>& reportedOutputs)
  {
      return getSubgraphTensors(newLayers, externalOutputs, true, &reportedOutputs);
  }
  
  Status getSubgraphInputs(const std::vector<nvinfer1::ILayer*>& newLayers,
      std::unordered_map<std::string, std::set<int32_t>>& externalInputs)
  {
      return getSubgraphTensors(newLayers, externalInputs, false);
  }
  
  } // namespace onnx2trt
  