get next gates up to max depth

include/hal_core/netlist/decorators/netlist_traversal_decorator.h

@@ -124,6 +124,38 @@ namespace hal
                                                               const std::function<bool(const Endpoint*, const u32 current_depth)>& exit_endpoint_filter  = nullptr,
                                                               const std::function<bool(const Endpoint*, const u32 current_depth)>& entry_endpoint_filter = nullptr) const;
 
+        /**
+         * Starting from the given net, traverse the netlist and return only the successor/predecessor gates for which the `target_gate_filter` evaluates to `true`.
+         * Continue traversal independent of whatever `target_gate_filter` evaluates to.
+         * Stop traversal if the specified depth is reached.
+         * The current depth is counted starting at 1 for the destinations of the provided net. 
+         * If no depth is provided, all gates between the start net and the global netlist outputs will be traversed.
+         * The target_gate_filter may be omitted in which case all traversed gates will be returned.
+         * 
+         * @param[in] net - Start net.
+         * @param[in] successors - Set `true` to get successors, set `false` to get predecessors.
+         * @param[in] target_gate_filter - Filter condition that must be met for the target gates.
+         * @param[in] max_depth - The maximum depth for netlist traversal starting from the start net.
+         * @returns The next gates fulfilling the target gate filter condition on success, an error otherwise.
+         */
+        Result<std::set<Gate*>> get_next_matching_gates_until_depth(const Net* net, bool successors, const std::function<bool(const Gate*)>& target_gate_filter = nullptr, u32 max_depth = 0) const;
+
+        /**
+         * Starting from the given gate, traverse the netlist and return only the successor/predecessor gates for which the `target_gate_filter` evaluates to `true`.
+         * Continue traversal independent of whatever `target_gate_filter` evaluates to.
+         * Stop traversal if the specified depth is reached.
+         * The current depth is counted starting at 1 for the direct successors/predecessors of the provided gate. 
+         * If no depth is provided, all gates between the start gate and the global netlist outputs will be traversed.
+         * The target_gate_filter may be omitted in which case all traversed gates will be returned.
+         * 
+         * @param[in] gate - Start gate.
+         * @param[in] successors - Set `true` to get successors, set `false` to get predecessors.
+         * @param[in] target_gate_filter - Filter condition that must be met for the target gates.
+         * @param[in] max_depth - The maximum depth for netlist traversal starting from the start gate.
+         * @returns The next gates fulfilling the target gate filter condition on success, an error otherwise.
+         */
+        Result<std::set<Gate*>> get_next_matching_gates_until_depth(const Gate* gate, bool successors, const std::function<bool(const Gate*)>& target_gate_filter = nullptr, u32 max_depth = 0) const;
+
         /**
          * Starting from the given net, traverse the netlist and return only the next layer of sequential successor/predecessor gates.
          * Traverse over gates that are not sequential until a sequential gate is found.

src/netlist/decorators/netlist_traversal_decorator.cpp

@@ -235,6 +235,108 @@ namespace hal
         return OK(res);
     }
 
+    Result<std::set<Gate*>>
+        NetlistTraversalDecorator::get_next_matching_gates_until_depth(const Net* net, bool successors, const std::function<bool(const Gate*)>& target_gate_filter, u32 max_depth) const
+    {
+        if (net == nullptr)
+        {
+            return ERR("nullptr given as net");
+        }
+
+        if (!m_netlist.is_net_in_netlist(net))
+        {
+            return ERR("net does not belong to netlist");
+        }
+
+        std::unordered_set<const Net*> visited;
+        std::vector<const Net*> stack = {net};
+        std::vector<const Net*> previous;
+        std::set<Gate*> res;
+        while (!stack.empty())
+        {
+            const Net* current = stack.back();
+
+            if (!previous.empty() && current == previous.back())
+            {
+                stack.pop_back();
+                previous.pop_back();
+                continue;
+            }
+
+            u32 current_depth = previous.size() + 1;
+            visited.insert(current);
+
+            bool added = false;
+            for (const auto* entry_ep : successors ? current->get_destinations() : current->get_sources())
+            {
+                if (max_depth != 0 && current_depth > max_depth)
+                {
+                    continue;
+                }
+
+                auto* g = entry_ep->get_gate();
+
+                if ((target_gate_filter == nullptr) || target_gate_filter(g))
+                {
+                    res.insert(g);
+                }
+
+                for (const auto* exit_ep : successors ? g->get_fan_out_endpoints() : g->get_fan_in_endpoints())
+                {
+                    if (max_depth != 0 && current_depth == max_depth)
+                    {
+                        continue;
+                    }
+
+                    const Net* n = exit_ep->get_net();
+                    if (visited.find(n) == visited.end())
+                    {
+                        stack.push_back(n);
+                        added = true;
+                    }
+                }
+            }
+
+            if (added)
+            {
+                previous.push_back(current);
+            }
+            else
+            {
+                stack.pop_back();
+            }
+        }
+
+        return OK(res);
+    }
+
+    Result<std::set<Gate*>>
+        NetlistTraversalDecorator::get_next_matching_gates_until_depth(const Gate* gate, bool successors, const std::function<bool(const Gate*)>& target_gate_filter, u32 max_depth) const
+    {
+        if (gate == nullptr)
+        {
+            return ERR("nullptr given as gate");
+        }
+
+        if (!m_netlist.is_gate_in_netlist(gate))
+        {
+            return ERR("net does not belong to netlist");
+        }
+
+        std::set<Gate*> res;
+        for (const auto* exit_ep : successors ? gate->get_fan_out_endpoints() : gate->get_fan_in_endpoints())
+        {
+            const auto next_res = this->get_next_matching_gates_until_depth(exit_ep->get_net(), successors, target_gate_filter, max_depth);
+            if (next_res.is_error())
+            {
+                return ERR(next_res.get_error());
+            }
+            auto next = next_res.get();
+            res.insert(next.begin(), next.end());
+        }
+        return OK(res);
+    }
+
     Result<std::set<Gate*>>
         NetlistTraversalDecorator::get_next_sequential_gates(const Net* net, bool successors, const std::set<PinType>& forbidden_pins, std::unordered_map<const Net*, std::set<Gate*>>* cache) const
     {

src/python_bindings/bindings/netlist_traversal_decorator.cpp

@@ -112,7 +112,7 @@ namespace hal
             },
             py::arg("net"),
             py::arg("successors"),
-            py::arg("target_gate_filter"),
+            py::arg("target_gate_filter")    = nullptr,
             py::arg("exit_endpoint_filter")  = nullptr,
             py::arg("entry_endpoint_filter") = nullptr,
             R"(
@@ -152,7 +152,7 @@ namespace hal
             },
             py::arg("gate"),
             py::arg("successors"),
-            py::arg("target_gate_filter"),
+            py::arg("target_gate_filter")    = nullptr,
             py::arg("exit_endpoint_filter")  = nullptr,
             py::arg("entry_endpoint_filter") = nullptr,
             R"(
@@ -171,6 +171,76 @@ namespace hal
             :rtype: set[hal_py.Gate] or None
         )");
 
+        py_netlist_traversal_decorator.def(
+            "get_next_matching_gates_until_depth",
+            [](NetlistTraversalDecorator& self, const Net* net, bool successors, const std::function<bool(const Gate*)>& target_gate_filter = nullptr, u32 depth = 0)
+                -> std::optional<std::set<Gate*>> {
+                auto res = self.get_next_matching_gates_until_depth(net, successors, target_gate_filter, depth);
+                if (res.is_ok())
+                {
+                    return res.get();
+                }
+                else
+                {
+                    log_error("python_context", "error encountered while getting next gates:\n{}", res.get_error().get());
+                    return std::nullopt;
+                }
+            },
+            py::arg("net"),
+            py::arg("successors"),
+            py::arg("target_gate_filter") = nullptr,
+            py::arg("depth")              = 0,
+            R"(
+            Starting from the given net, traverse the netlist and return only the successor/predecessor gates for which the ``target_gate_filter`` evaluates to ``True``.
+            Continue traversal independent of whatever ``target_gate_filter`` evaluates to.
+            Stop traversal if the specified depth is reached.
+            The current depth is counted starting at 1 for the destinations of the provided net. 
+            If no depth is provided, all gates between the start net and the global netlist outputs will be traversed.
+            The target_gate_filter may be omitted in which case all traversed gates will be returned.
+
+            :param hal_py.Net net: Start net.
+            :param bool successors: Set ``True`` to get successors, set ``False`` to get predecessors.
+            :param lambda target_gate_filter: Filter condition that must be met for the target gates.
+            :param int depth: The maximum depth for netlist traversal starting from the start net.
+            :returns: The next gates fulfilling the target gate filter condition on success, ``None`` otherwise.
+            :rtype: set[hal_py.Gate] or None
+        )");
+
+        py_netlist_traversal_decorator.def(
+            "get_next_matching_gates_until_depth",
+            [](NetlistTraversalDecorator& self, const Gate* gate, bool successors, const std::function<bool(const Gate*)>& target_gate_filter = nullptr, u32 depth = 0)
+                -> std::optional<std::set<Gate*>> {
+                auto res = self.get_next_matching_gates_until_depth(gate, successors, target_gate_filter, depth);
+                if (res.is_ok())
+                {
+                    return res.get();
+                }
+                else
+                {
+                    log_error("python_context", "error encountered while getting next gates:\n{}", res.get_error().get());
+                    return std::nullopt;
+                }
+            },
+            py::arg("gate"),
+            py::arg("successors"),
+            py::arg("target_gate_filter") = nullptr,
+            py::arg("depth")              = 0,
+            R"(
+            Starting from the given gate, traverse the netlist and return only the successor/predecessor gates for which the ``target_gate_filter`` evaluates to ``True``.
+            Continue traversal independent of whatever ``target_gate_filter`` evaluates to.
+            Stop traversal if the specified depth is reached.
+            The current depth is counted starting at 1 for the direct successors/predecessors of the provided gate. 
+            If no depth is provided, all gates between the start gate and the global netlist outputs will be traversed.
+            The target_gate_filter may be omitted in which case all traversed gates will be returned.
+
+            :param hal_py.Gate gate: Start gate.
+            :param bool successors: Set ``True`` to get successors, set ``False`` to get predecessors.
+            :param lambda target_gate_filter: Filter condition that must be met for the target gates.
+            :param int depth: The maximum depth for netlist traversal starting from the start gate.
+            :returns: The next gates fulfilling the target gate filter condition on success, ``None`` otherwise.
+            :rtype: set[hal_py.Gate] or None
+        )");
+
         py_netlist_traversal_decorator.def(
             "get_next_sequential_gates",
             [](NetlistTraversalDecorator& self, const Net* net, bool successors, const std::set<PinType>& forbidden_pins) -> std::optional<std::set<Gate*>> {