Update AMT max index bound (#714)

ipld/amt/src/amt.rs

@@ -1,7 +1,9 @@
 // Copyright 2020 ChainSafe Systems
 // SPDX-License-Identifier: Apache-2.0, MIT
 
-use crate::{node::Link, nodes_for_height, BitMap, Error, Node, Root, MAX_INDEX, WIDTH};
+use crate::{
+    node::Link, nodes_for_height, BitMap, Error, Node, Root, MAX_HEIGHT, MAX_INDEX, WIDTH,
+};
 use cid::{multihash::Blake2b256, Cid};
 use encoding::{de::DeserializeOwned, ser::Serialize};
 use ipld_blockstore::BlockStore;
@@ -58,11 +60,16 @@ where
             .get(cid)?
             .ok_or_else(|| Error::CidNotFound(cid.to_string()))?;
 
+        // Sanity check, this should never be possible.
+        if root.height > MAX_HEIGHT {
+            return Err(Error::MaxHeight(root.height, MAX_HEIGHT));
+        }
+
         Ok(Self { root, block_store })
     }
 
     // Getter for height
-    pub fn height(&self) -> u32 {
+    pub fn height(&self) -> u64 {
         self.root.height
     }
 
@@ -82,7 +89,7 @@ where
 
     /// Get value at index of AMT
     pub fn get(&self, i: u64) -> Result<Option<V>, Error> {
-        if i >= MAX_INDEX {
+        if i > MAX_INDEX {
             return Err(Error::OutOfRange(i));
         }
 
@@ -95,11 +102,11 @@ where
 
     /// Set value at index
     pub fn set(&mut self, i: u64, val: V) -> Result<(), Error> {
-        if i >= MAX_INDEX {
+        if i > MAX_INDEX {
             return Err(Error::OutOfRange(i));
         }
 
-        while i >= nodes_for_height(self.height() + 1 as u32) {
+        while i >= nodes_for_height(self.height() + 1) {
             // node at index exists
             if !self.root.node.empty() {
                 // Save and get cid to be able to link from higher level node
@@ -150,7 +157,7 @@ where
 
     /// Delete item from AMT at index
     pub fn delete(&mut self, i: u64) -> Result<bool, Error> {
-        if i >= MAX_INDEX {
+        if i > MAX_INDEX {
             return Err(Error::OutOfRange(i));
         }
 

ipld/amt/src/error.rs

@@ -13,6 +13,9 @@ pub enum Error {
     /// Index referenced it above arbitrary max set
     #[error("index {0} out of range for the amt")]
     OutOfRange(u64),
+    /// Height of root node is greater than max.
+    #[error("failed to load AMT: height out of bounds: {0} > {1}")]
+    MaxHeight(u64, u64),
     /// Cbor encoding error
     #[error(transparent)]
     Encoding(#[from] EncodingError),

ipld/amt/src/lib.rs

@@ -18,9 +18,21 @@ pub use self::error::Error;
 pub(crate) use self::node::Node;
 pub(crate) use self::root::Root;
 
-const WIDTH: usize = 8;
-pub const MAX_INDEX: u64 = 1 << 48;
+const MAX_INDEX_BITS: u64 = 63;
+const WIDTH_BITS: u64 = 3;
+const WIDTH: usize = 1 << WIDTH_BITS; // 8
+const MAX_HEIGHT: u64 = MAX_INDEX_BITS / WIDTH_BITS - 1;
 
-pub(crate) fn nodes_for_height(height: u32) -> u64 {
-    (WIDTH as u64).pow(height)
+// Maximum index for elements in the AMT. This is currently 1^63
+// (max int) because the width is 8. That means every "level" consumes 3 bits
+// from the index, and 63/3 is a nice even 21
+pub const MAX_INDEX: u64 = (1 << MAX_INDEX_BITS) - 1;
+
+fn nodes_for_height(height: u64) -> u64 {
+    let height_log_two = WIDTH_BITS * height;
+    assert!(
+        height_log_two < 64,
+        "height overflow, should be checked at all entry points"
+    );
+    1 << height_log_two
 }

ipld/amt/src/node.rs

@@ -175,7 +175,7 @@ where
     pub(super) fn get<DB: BlockStore>(
         &self,
         bs: &DB,
-        height: u32,
+        height: u64,
         i: u64,
     ) -> Result<Option<V>, Error> {
         let sub_i = i / nodes_for_height(height);
@@ -204,7 +204,7 @@ where
     pub(super) fn set<DB: BlockStore>(
         &mut self,
         bs: &DB,
-        height: u32,
+        height: u64,
         i: u64,
         val: V,
     ) -> Result<bool, Error> {
@@ -269,7 +269,7 @@ where
     pub(super) fn delete<DB: BlockStore>(
         &mut self,
         bs: &DB,
-        height: u32,
+        height: u64,
         i: u64,
     ) -> Result<bool, Error> {
         let sub_i = i / nodes_for_height(height);
@@ -319,7 +319,7 @@ where
     pub(super) fn for_each<S, F>(
         &self,
         store: &S,
-        height: u32,
+        height: u64,
         offset: u64,
         f: &mut F,
     ) -> Result<(), Box<dyn StdError>>

ipld/amt/src/root.rs

@@ -10,7 +10,7 @@ use encoding::{
 /// Root of an AMT vector, can be serialized and keeps track of height and count
 #[derive(PartialEq, Debug)]
 pub(super) struct Root<V> {
-    pub(super) height: u32,
+    pub(super) height: u64,
     pub(super) count: u64,
     pub(super) node: Node<V>,
 }

ipld/amt/tests/amt_tests.rs

@@ -49,15 +49,15 @@ fn out_of_range() {
     let db = db::MemoryDB::default();
     let mut a = Amt::new(&db);
 
-    let res = a.set(1 << 50, "test".to_owned());
-    assert_eq!(res.err(), Some(Error::OutOfRange(1 << 50)));
+    let res = a.set((1 << 63) + 4, "test".to_owned());
+    assert!(matches!(res, Err(Error::OutOfRange(_))));
 
-    let res = a.set(MAX_INDEX, "test".to_owned());
-    assert_eq!(res.err(), Some(Error::OutOfRange(MAX_INDEX)));
+    let res = a.set(MAX_INDEX + 1, "test".to_owned());
+    assert!(matches!(res, Err(Error::OutOfRange(_))));
 
-    let res = a.set(MAX_INDEX - 1, "test".to_owned());
+    let res = a.set(MAX_INDEX, "test".to_owned());
     assert_eq!(res.err(), None);
-    assert_get(&mut a, MAX_INDEX - 1, &"test".to_owned());
+    assert_get(&mut a, MAX_INDEX, &"test".to_owned());
 }
 
 #[test]