Skip to content

Commit

Permalink
Fix memory allocation issues in binomial graph generators
Browse files Browse the repository at this point in the history 
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 Qiskit#457
  • Loading branch information
@mtreinish
mtreinish committed Oct 2, 2021
1 parent 6d62503 commit 73e8968
Show file tree
Hide file tree
Showing 3 changed files with 90 additions and 92 deletions.
11 changes: 11 additions & 0 deletions releasenotes/notes/binomial-tree-alloc-fix-c24c8c4f3c27d489.yaml
View file
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?>`__
163 changes: 71 additions & 92 deletions src/generators.rs
View file
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 @@ -884,59 +884,51 @@ pub fn binomial_tree_graph(
weights: Option<Vec<PyObject>>,
multigraph: bool,
) -> PyResult<graph::PyGraph> {
let mut graph = StableUnGraph::<PyObject, PyObject>::default();

if order >= 60 {
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 Down Expand Up @@ -984,39 +976,38 @@ pub fn directed_binomial_tree_graph(
bidirectional: bool,
multigraph: bool,
) -> PyResult<digraph::PyDiGraph> {
let mut graph = StableDiGraph::<PyObject, PyObject>::default();

if order >= 60 {
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 +1016,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
View file
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 73e8968

Please sign in to comment.