From ceb1018ba7717c3ef0a7a5ac624eb16652d90c08 Mon Sep 17 00:00:00 2001
From: Lachlan Deakin <ljdgit@gmail.com>
Date: Sun, 11 Feb 2024 18:04:39 +1100
Subject: [PATCH] Add `ChunkShape::num_elements`

---
 CHANGELOG.md             |  1 +
 src/array/chunk_shape.rs | 13 +++++++++++++
 2 files changed, 14 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 358182ff..4dc4efe1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -19,6 +19,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
  - Additional sharding tests
  - Add store lock tests
  - Added `contiguous_elements` method to `ContiguousIndicesIterator` and `ContiguousLinearisedIndicesIterator`
+ - Added `ChunkShape::num_elements`
 
 ### Changed
  - Dependency bumps
diff --git a/src/array/chunk_shape.rs b/src/array/chunk_shape.rs
index 0b26d507..bb914d62 100644
--- a/src/array/chunk_shape.rs
+++ b/src/array/chunk_shape.rs
@@ -5,9 +5,22 @@ use serde::{Deserialize, Serialize};
 use super::{ArrayShape, NonZeroError};
 
 /// The shape of a chunk. All dimensions must be non-zero.
+#[allow(clippy::unsafe_derive_deserialize)]
 #[derive(Serialize, Deserialize, Clone, Eq, PartialEq, Debug)]
 pub struct ChunkShape(Vec<NonZeroU64>);
 
+impl ChunkShape {
+    /// Return the number of elements.
+    ///
+    /// Equal to the product of the components of its shape.
+    pub fn num_elements(&self) -> NonZeroU64 {
+        unsafe {
+            // Multiplying NonZeroU64 must result in NonZeroU64
+            NonZeroU64::new_unchecked(self.0.iter().copied().map(NonZeroU64::get).product::<u64>())
+        }
+    }
+}
+
 impl From<ChunkShape> for Vec<NonZeroU64> {
     fn from(val: ChunkShape) -> Self {
         val.0.clone()