From ec97d1aa80deba74476d925fa2fb3410d69e2ff9 Mon Sep 17 00:00:00 2001
From: firejoust <42525540+firejoust@users.noreply.github.com>
Date: Wed, 30 Nov 2022 15:36:14 +1100
Subject: [PATCH] Get blocks by state ID

---
 src/api/blocks.rs       | 24 +++++++++++++++++++-----
 src/api/tests/blocks.rs | 12 ++++++++++++
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/src/api/blocks.rs b/src/api/blocks.rs
index 2b14dc6..373584c 100644
--- a/src/api/blocks.rs
+++ b/src/api/blocks.rs
@@ -23,11 +23,17 @@ impl Blocks {
         Ok(blocks)
     }
 
-    /// Returns the blocks indexed by ID
-    pub fn blocks(&self) -> DataResult<HashMap<u32, Block>> {
-        let blocks = self.blocks_array()?;
-        let blocks_map = blocks.into_iter().map(|b| (b.id, b)).collect();
-
+    // Returns the references of blocks indexed by state ID
+    pub fn blocks_by_state_id<'a>(&self, blocks: &'a Vec<Block>) -> DataResult<HashMap<u32, &'a Block>> {
+        let mut blocks_map = HashMap::new();
+        blocks.iter().for_each(|b| {
+            let min_state_id = b.min_state_id.unwrap_or(b.id << 4);
+            let max_state_id = b.max_state_id.unwrap_or(min_state_id + 15);
+            (min_state_id..max_state_id).for_each(|s| {
+                blocks_map.insert(s, b);
+            });
+        });
+        
         Ok(blocks_map)
     }
 
@@ -39,6 +45,14 @@ impl Blocks {
         Ok(blocks_map)
     }
 
+    /// Returns the blocks indexed by ID
+    pub fn blocks(&self) -> DataResult<HashMap<u32, Block>> {
+        let blocks = self.blocks_array()?;
+        let blocks_map = blocks.into_iter().map(|b| (b.id, b)).collect();
+
+        Ok(blocks_map)
+    }
+
     /// Returns the block collision shapes object
     pub fn block_collision_shapes(&self) -> DataResult<BlockCollisionShapes> {
         let content = get_version_specific_file(&self.version, BLOCK_COLLISION_SHAPES_FILE)?;
diff --git a/src/api/tests/blocks.rs b/src/api/tests/blocks.rs
index 5570dbe..fb27a91 100644
--- a/src/api/tests/blocks.rs
+++ b/src/api/tests/blocks.rs
@@ -11,6 +11,18 @@ pub fn test_blocks_array() {
     }
 }
 
+#[test]
+pub fn test_blocks_by_state_id() {
+    let versions = get_test_versions();
+
+    for version in versions {
+        let api = get_api(version);
+        let blocks = api.blocks.blocks_array().unwrap();
+        let by_state = api.blocks.blocks_by_state_id(&blocks);
+        assert!(by_state.is_ok());
+    }
+}
+
 #[test]
 pub fn test_blocks_by_name() {
     let versions = get_test_versions();