//*****************************************************************************
// Copyright 2017-2021 Intel Corporation
//
// Licensed 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.
//*****************************************************************************

#include "pass/shape_relevance.hpp"
#include "ngraph/graph_util.hpp"
#include "ngraph/op/constant.hpp"
#include "ngraph/op/util/op_types.hpp"

using namespace ngraph;

//
// This pass refreshes the "is_relevant_to_shape" flag on each parameter. A parameter will be
// flagged as relevant to shapes if there is any path from that parameter to a shape-relevant
// input that does _not_ pass through a value-irrelevant input. For example:
//
//      N0[Parameter] N1[Parameter]
//       |            |
//       |            |
//       |            |
//       N2[v1::Reshape]
//
// N1 (but not N0) will be flagged as shape-relevant, because N1 feeds into the "shape" input
// of N2.
//
//      N0[Parameter] N1[Parameter]
//       |            |
//       |            N2[ShapeOf]
//       |            |
//       N3[v1::Reshape]
//
// Neither N0 nor N1 will be flagged as shape-relevant. (N1 does feed into the "shape" input of N3,
// but only via the value-irrelevant input of ShapeOf.)
//
bool pass::ShapeRelevance::run_on_function(std::shared_ptr<Function> f)
{
    // TODO(amprocte): We are probably reinventing the wheel with the graph traversal here; the
    // reason is that we need to cut the traversal short in cases where input values are
    // irrelevant. See if there is a way to reduce this duplication.

    // Set of nodes that must be evaluated to determine the value of shape-relevant inputs.
    std::set<Node*> shape_determinants;

    // Step 1: Find root nodes (these are nodes with an output connected to a shape-relevant
    // input).
    for (auto& n : f->get_ops())
    {
        for (auto& output : n->outputs())
        {
            for (auto& input : output.get_target_inputs())
            {
                if (input.get_is_relevant_to_shapes())
                {
                    shape_determinants.insert(n.get());
                    break;
                }
            }
        }
    }

    // Step 2: Find all shape determinants. This is the transitive closure of R, where n1 R n2
    // iff there is a data flow edge from n2 to n1 and that data flow edge is not
    // value-irrelevant.
    bool changes_made = false;

    {
        std::list<Node*> to_visit{shape_determinants.begin(), shape_determinants.end()};
        std::set<Node*> already_visited;

        while (!to_visit.empty())
        {
            auto node = to_visit.front();
            to_visit.pop_front();

            if (already_visited.count(node) > 0)
            {
                continue;
            }

            shape_determinants.insert(node);
            already_visited.insert(node);

            if (op::is_parameter(node))
            {
                auto node_as_param = static_cast<op::Parameter*>(node);
                if (!node_as_param->is_relevant_to_shapes())
                {
                    node_as_param->set_is_relevant_to_shapes(true);
                    changes_made = true;
                }
            }

            for (size_t i = 0; i < node->get_input_size(); i++)
            {
                if (!node->input(i).get_is_relevant_to_values())
                {
                    continue;
                }
                auto source_node = node->get_input_node_ptr(i);
                if (already_visited.count(source_node) == 0)
                {
                    to_visit.push_front(source_node);
                }
            }
        }
    }

    return changes_made;
}
