Skip to content

Commit

Permalink
Fix memory allocation issues in binomial graph generators (#458)
Browse files Browse the repository at this point in the history
* Fix memory allocation issues in binomial graph generators

This commit fixes issues with the allocation of memory in the binomial
graph generators. Previously there were 2 issues in the binomial graph
generator functions when it came to allocating memory. The first was a
temporary Vec was used to store all the node indices in the graph which
would eat num_nodes * sizeof(usize) bytes of memory. This Vec isn't
actually needed as the NodeIndex is just a usize integer wrapper and we
know the indices of the nodes. This commit removes this and just
iterates over the range of nodes instead of creating a Vec reducing the
amoount of memory required. The second issue was that it was easily
possible to overflow the max Vec size (either inside the graph objects
or from the Vec previously created to store the node indices) by using
an order where 2**order > isize::MAX or if order > 64 causing an
overflow of a 64bit usize which would result in a panic. This fixes
this issue by checking the input of order and seeing if it will overflow
or exceed the max Vec size and raise an OverflowError.

Fixes #457

* Use a constant value dependent on platform pointer width

* Update src/generators.rs

Co-authored-by: Ivan Carvalho <8753214+IvanIsCoding@users.noreply.github.com>
  • Loading branch information
mtreinish and IvanIsCoding committed Oct 10, 2021
1 parent 6d62503 commit 0d7fdf4
Show file tree
Hide file tree
Showing 3 changed files with 110 additions and 95 deletions.
11 changes: 11 additions & 0 deletions releasenotes/notes/binomial-tree-alloc-fix-c24c8c4f3c27d489.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
---
fixes:
- |
Fixed an issue with the :func:`~retworkx.generators.binomial_tree_graph`
and :func:`~retworkx.generators.directed_binomial_tree_graph` generator
functions in :mod:`retworkx.generators` where passing an ``order`` value
``>= 60`` would cause an overflow and raise a ``PanicException`` caused by
the internal Rust panic when overflowing or exceeding the max Vec size.
Instead the function will raise an ``OverflowError`` and indicate the
specified order is too large.
Fixed `#457 <https://github.com/Qiskit/retworkx/issues/457?>`__
186 changes: 91 additions & 95 deletions src/generators.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ use petgraph::graph::NodeIndex;
use petgraph::stable_graph::{StableDiGraph, StableUnGraph};
use petgraph::visit::{EdgeRef, IntoEdgeReferences};

use pyo3::exceptions::PyIndexError;
use pyo3::exceptions::{PyIndexError, PyOverflowError};
use pyo3::prelude::*;
use pyo3::wrap_pyfunction;
use pyo3::Python;
Expand Down Expand Up @@ -854,9 +854,19 @@ pub fn directed_grid_graph(
})
}

// MAX_ORDER is determined based on the pointer width of the target platform
#[cfg(target_pointer_width = "64")]
const MAX_ORDER: u32 = 60;
#[cfg(not(target_pointer_width = "64"))]
const MAX_ORDER: u32 = 29;

/// Generate an undirected binomial tree of order n recursively.
///
/// :param int order: Order of the binomial tree.
/// :param int order: Order of the binomial tree. The maximum allowed value
/// for order on the platform your running on. If it's a 64bit platform
/// the max value is 59 and on 32bit systems the max value is 29. Any order
/// value above these will raise a ``OverflowError``.
/// depends
/// :param list weights: A list of node weights. If the number of weights is
/// less than 2**order extra nodes with with None will be appended.
/// :param bool multigraph: When set to False the output
Expand All @@ -866,7 +876,9 @@ pub fn directed_grid_graph(
///
/// :returns: A binomial tree with 2^n vertices and 2^n - 1 edges.
/// :rtype: PyGraph
/// :raises IndexError: If the lenght of ``weights`` is greater that 2^n
/// :raises IndexError: If the length of ``weights`` is greater that 2^n
/// :raises OverflowError: If the input order exceeds the maximum value for the
/// current platform.
///
/// .. jupyter-execute::
///
Expand All @@ -884,59 +896,51 @@ pub fn binomial_tree_graph(
weights: Option<Vec<PyObject>>,
multigraph: bool,
) -> PyResult<graph::PyGraph> {
let mut graph = StableUnGraph::<PyObject, PyObject>::default();

if order >= MAX_ORDER {
return Err(PyOverflowError::new_err(format!(
"An order of {} exceeds the max allowable size",
order
)));
}
let num_nodes = usize::pow(2, order);

let nodes: Vec<NodeIndex> = match weights {
Some(weights) => {
let mut node_list: Vec<NodeIndex> = Vec::new();

let mut node_count = num_nodes;

if weights.len() > num_nodes {
return Err(PyIndexError::new_err(
"weights should be <= 2**order",
));
}

for weight in weights {
let index = graph.add_node(weight);
node_list.push(index);
node_count -= 1;
}

for _i in 0..node_count {
let index = graph.add_node(py.None());
node_list.push(index);
let num_edges = usize::pow(2, order) - 1;
let mut graph = StableUnGraph::<PyObject, PyObject>::with_capacity(
num_nodes, num_edges,
);
for i in 0..num_nodes {
match weights {
Some(ref weights) => {
if weights.len() > num_nodes {
return Err(PyIndexError::new_err(
"weights should be <= 2**order",
));
}
if i < weights.len() {
graph.add_node(weights[i].clone_ref(py))
} else {
graph.add_node(py.None())
}
}

node_list
}

None => (0..num_nodes).map(|_| graph.add_node(py.None())).collect(),
};
None => graph.add_node(py.None()),
};
}

let mut n = 1;
let zero_index = NodeIndex::new(0);

for _ in 0..order {
let edges: Vec<(NodeIndex, NodeIndex)> = graph
.edge_references()
.map(|e| (e.source(), e.target()))
.collect();

for (source, target) in edges {
let source_index = source.index();
let target_index = target.index();
let source_index = NodeIndex::new(source.index() + n);
let target_index = NodeIndex::new(target.index() + n);

graph.add_edge(
nodes[source_index + n],
nodes[target_index + n],
py.None(),
);
graph.add_edge(source_index, target_index, py.None());
}

graph.add_edge(nodes[0], nodes[n], py.None());
graph.add_edge(zero_index, NodeIndex::new(n), py.None());

n *= 2;
}
Expand All @@ -951,7 +955,10 @@ pub fn binomial_tree_graph(
/// Generate an undirected binomial tree of order n recursively.
/// The edges propagate towards right and bottom direction if ``bidirectional`` is ``false``
///
/// :param int order: Order of the binomial tree.
/// :param int order: Order of the binomial tree. The maximum allowed value
/// for order on the platform your running on. If it's a 64bit platform
/// the max value is 59 and on 32bit systems the max value is 29. Any order
/// value above these will raise a ``OverflowError``.
/// :param list weights: A list of node weights. If the number of weights is
/// less than 2**order extra nodes with None will be appended.
/// :param bidirectional: A parameter to indicate if edges should exist in
Expand All @@ -964,6 +971,8 @@ pub fn binomial_tree_graph(
/// :returns: A directed binomial tree with 2^n vertices and 2^n - 1 edges.
/// :rtype: PyDiGraph
/// :raises IndexError: If the lenght of ``weights`` is greater that 2^n
/// :raises OverflowError: If the input order exceeds the maximum value for the
/// current platform.
///
/// .. jupyter-execute::
///
Expand All @@ -984,39 +993,38 @@ pub fn directed_binomial_tree_graph(
bidirectional: bool,
multigraph: bool,
) -> PyResult<digraph::PyDiGraph> {
let mut graph = StableDiGraph::<PyObject, PyObject>::default();

if order >= MAX_ORDER {
return Err(PyOverflowError::new_err(format!(
"An order of {} exceeds the max allowable size",
order
)));
}
let num_nodes = usize::pow(2, order);

let nodes: Vec<NodeIndex> = match weights {
Some(weights) => {
let mut node_list: Vec<NodeIndex> = Vec::new();
let mut node_count = num_nodes;

if weights.len() > num_nodes {
return Err(PyIndexError::new_err(
"weights should be <= 2**order",
));
}

for weight in weights {
let index = graph.add_node(weight);
node_list.push(index);
node_count -= 1;
}

for _i in 0..node_count {
let index = graph.add_node(py.None());
node_list.push(index);
let num_edges = usize::pow(2, order) - 1;
let mut graph = StableDiGraph::<PyObject, PyObject>::with_capacity(
num_nodes, num_edges,
);

for i in 0..num_nodes {
match weights {
Some(ref weights) => {
if weights.len() > num_nodes {
return Err(PyIndexError::new_err(
"weights should be <= 2**order",
));
}
if i < weights.len() {
graph.add_node(weights[i].clone_ref(py))
} else {
graph.add_node(py.None())
}
}

node_list
}

None => (0..num_nodes).map(|_| graph.add_node(py.None())).collect(),
};
None => graph.add_node(py.None()),
};
}

let mut n = 1;
let zero_index = NodeIndex::new(0);

for _ in 0..order {
let edges: Vec<(NodeIndex, NodeIndex)> = graph
Expand All @@ -1025,39 +1033,27 @@ pub fn directed_binomial_tree_graph(
.collect();

for (source, target) in edges {
let source_index = source.index();
let target_index = target.index();
let source_index = NodeIndex::new(source.index() + n);
let target_index = NodeIndex::new(target.index() + n);

if graph
.find_edge(nodes[source_index + n], nodes[target_index + n])
.is_none()
{
graph.add_edge(
nodes[source_index + n],
nodes[target_index + n],
py.None(),
);
if graph.find_edge(source_index, target_index).is_none() {
graph.add_edge(source_index, target_index, py.None());
}

if bidirectional
&& graph
.find_edge(nodes[target_index + n], nodes[source_index + n])
.is_none()
&& graph.find_edge(target_index, source_index).is_none()
{
graph.add_edge(
nodes[target_index + n],
nodes[source_index + n],
py.None(),
);
graph.add_edge(target_index, source_index, py.None());
}
}
let n_index = NodeIndex::new(n);

if graph.find_edge(nodes[0], nodes[n]).is_none() {
graph.add_edge(nodes[0], nodes[n], py.None());
if graph.find_edge(zero_index, n_index).is_none() {
graph.add_edge(zero_index, n_index, py.None());
}

if bidirectional && graph.find_edge(nodes[n], nodes[0]).is_none() {
graph.add_edge(nodes[n], nodes[0], py.None());
if bidirectional && graph.find_edge(n_index, zero_index).is_none() {
graph.add_edge(n_index, zero_index, py.None());
}

n *= 2;
Expand Down
8 changes: 8 additions & 0 deletions tests/generators/test_binomial_tree.py
Original file line number Diff line number Diff line change
Expand Up @@ -199,3 +199,11 @@ def test_directed_binomial_tree_graph_bidirectional(self):
self.assertEqual(len(graph), 2 ** n)
self.assertEqual(len(graph.edges()), 2 * (2 ** n - 1))
self.assertEqual(list(graph.edge_list()), expected_edges[n])

def test_overflow_binomial_tree(self):
with self.assertRaises(OverflowError):
retworkx.generators.binomial_tree_graph(75)

def test_overflow_directed_binomial_tree(self):
with self.assertRaises(OverflowError):
retworkx.generators.directed_binomial_tree_graph(75)

0 comments on commit 0d7fdf4

Please sign in to comment.