From ce1cc868f16f75a01b29374416616c0e1a612b2e Mon Sep 17 00:00:00 2001
From: Dung Vo <dungv@google.com>
Date: Wed, 28 Oct 2026 06:49:27 -0700
Subject: [PATCH] lzma_rust2 binding: report memory usage for the stream, auto,
 alone and lzip decoders

`lzma_memusage` reported 0 for every decoder that learns its dictionary size from the input.
It now returns the crate's `memory_usage()`, the figure checked against the memory limit, once
the header that decides it has been parsed; the auto decoder answers for the format it
detected.

PiperOrigin-RevId: 1008872197
---
 lzma/lzma_rust_api.h |  6 +++---
 lzma/rust/state.rs   | 17 ++++++++++++-----
 lzma/rust/util.rs    |  8 ++++----
 3 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/lzma/lzma_rust_api.h b/lzma/lzma_rust_api.h
index 060fb43d..f8f4f2a6 100644
--- a/lzma/lzma_rust_api.h
+++ b/lzma/lzma_rust_api.h
@@ -546,10 +546,10 @@ lzma_check lzma_get_check(const lzma_stream* strm);
 // For encoders and raw decoders where options are explicitly configured at
 // initialization, returns the calculated memory usage in bytes.
 //
-// Note: For streaming decoders where dictionary size is determined dynamically
+// For streaming decoders where dictionary size is determined dynamically
 // from the stream header (e.g., `lzma_stream_decoder`, `lzma_alone_decoder`,
-// `lzma_auto_decoder`, `lzma_lzip_decoder`), this implementation returns 0
-// as the underlying Rust decoders do not expose runtime memory usage.
+// `lzma_auto_decoder`, `lzma_lzip_decoder`), returns the memory usage in bytes
+// once the header has been parsed, and 0 before that.
 uint64_t lzma_memusage(const lzma_stream* strm);
 
 #ifdef __cplusplus
diff --git a/lzma/rust/state.rs b/lzma/rust/state.rs
index 7dc5fc8a..d2409909 100644
--- a/lzma/rust/state.rs
+++ b/lzma/rust/state.rs
@@ -201,7 +201,12 @@ impl InternalState {
     }
 
     pub fn memusage(&self) -> u64 {
-        match &self.coder {
+        // The auto decoder answers for the format it detected.
+        let coder = match &self.coder {
+            CoderInner::AutoDecoder(auto) => auto.inner().unwrap_or(&self.coder),
+            coder => coder,
+        };
+        match coder {
             CoderInner::RawLzma2Decoder(_) => {
                 if let Some(opts) = &self.lzma_options {
                     let dict_size = opts.dict_size.max(lzma_rust2::DICT_SIZE_MIN);
@@ -209,7 +214,7 @@ impl InternalState {
                 }
                 0
             }
-            CoderInner::LzmaDecoder(_) => {
+            CoderInner::LzmaDecoder(lzma) => {
                 if let Some(opts) = &self.lzma_options {
                     let dict_size = opts.dict_size.max(lzma_rust2::DICT_SIZE_MIN);
                     if let Ok(mem_kb) =
@@ -218,10 +223,12 @@ impl InternalState {
                         return (mem_kb as u64) * 1024;
                     }
                 }
-                0
+                lzma.memory_usage().map_or(0, |kb| u64::from(kb) * 1024)
+            }
+            CoderInner::XzDecoder(xz) => xz.memory_usage().map_or(0, |kb| u64::from(kb) * 1024),
+            CoderInner::LzipDecoder(lzip) => {
+                lzip.memory_usage().map_or(0, |kb| u64::from(kb) * 1024)
             }
-            // NOTE (b/562586054): Add support for lzma_stream_decoder,
-            // lzma_auto_decoder, lzma_lzip_decoder, and lzma_alone_decoder.
             _ => 0,
         }
     }
diff --git a/lzma/rust/util.rs b/lzma/rust/util.rs
index 00db4888..9706bd7e 100644
--- a/lzma/rust/util.rs
+++ b/lzma/rust/util.rs
@@ -27,15 +27,15 @@ pub extern "C" fn lzma_cputhreads() -> u32 {
 ///
 /// Implements `lzma_memusage` by interrogating the rust state struct.
 /// Returns memory usage for encoders and raw decoders where options are known.
-/// For streaming decoders (XZ, Alone, Lzip, Auto), returns 0 because
-/// the underlying decoders do not expose dynamically parsed dictionary size.
+/// For streaming decoders (XZ, Alone, Lzip, Auto), returns the figure the crate
+/// checks against the memory limit, which is 0 until the header has been parsed.
 ///
 /// # Safety
 /// - `strm` must be null or point to a valid, properly aligned `lzma_stream`.
 /// - If `strm` is non-null and `strm.internal` is non-null, it must point to a valid
 ///   `InternalState` allocated by this crate.
-/// - The caller must guarantee that `*strm` and `*strm.internal` are not concurrently modified
-///   or deallocated during this call.
+/// - The caller must guarantee that `*strm` and `*strm.internal` are not concurrently modified or
+///   deallocated during this call.
 #[unsafe(export_name = crate::prefix!(lzma_memusage))]
 pub unsafe extern "C" fn lzma_memusage(strm: Option<&crate::types::lzma_stream>) -> u64 {
     let Some(strm) = strm else {