Changes the ParametersManager::Update*State()

- Validate on the NEXT start timestamp instead of the current one.
- Not use a callback function for additional updating steps.

PiperOrigin-RevId: 695203833

iamf/cli/BUILD

@@ -367,6 +367,7 @@ cc_library(
     hdrs = ["parameters_manager.h"],
     deps = [
         ":audio_element_with_data",
+        ":cli_util",
         ":parameter_block_with_data",
         "//iamf/common:macros",
         "//iamf/obu:demixing_param_definition",

iamf/cli/parameters_manager.cc

@@ -13,6 +13,7 @@
 
 #include <algorithm>
 #include <cstdint>
+#include <optional>
 
 #include "absl/container/flat_hash_map.h"
 #include "absl/log/check.h"
@@ -21,6 +22,7 @@
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "iamf/cli/audio_element_with_data.h"
+#include "iamf/cli/cli_util.h"
 #include "iamf/cli/parameter_block_with_data.h"
 #include "iamf/common/macros.h"
 #include "iamf/obu/demixing_info_parameter_data.h"
@@ -34,13 +36,14 @@ namespace iamf_tools {
 
 namespace {
 
-template <typename StateType, typename StateUpdater>
+template <typename StateType>
 absl::Status UpdateParameterState(
-    DecodedUleb128 audio_element_id, int32_t expected_timestamp,
+    DecodedUleb128 audio_element_id, int32_t expected_next_timestamp,
     absl::flat_hash_map<DecodedUleb128, StateType>& parameter_states,
     absl::flat_hash_map<DecodedUleb128, const ParameterBlockWithData*>&
         parameter_blocks,
-    absl::string_view parameter_name, StateUpdater additional_state_updater) {
+    absl::string_view parameter_name,
+    std::optional<StateType*>& parameter_state) {
   const auto parameter_states_iter = parameter_states.find(audio_element_id);
   if (parameter_states_iter == parameter_states.end()) {
     // No parameter definition found for the audio element ID, so
@@ -49,36 +52,33 @@ absl::Status UpdateParameterState(
   }
 
   // Validate the timestamps before updating.
-  auto& parameter_state = parameter_states_iter->second;
+  parameter_state = &parameter_states_iter->second;
 
   // Using `.at()` here is safe because if the parameter state exists for the
   // `audio_element_id`, an entry in `parameter_blocks` with the key
-  // `parameter_state.param_definition->parameter_id_` has already been
+  // `parameter_state->param_definition->parameter_id_` has already been
   // created during `Initialize()`.
   auto& parameter_block =
-      parameter_blocks.at(parameter_state.param_definition->parameter_id_);
+      parameter_blocks.at((*parameter_state)->param_definition->parameter_id_);
   if (parameter_block == nullptr) {
     // No parameter block found for this ID. Do not validate the timestamp
-    // or update anything else.
+    // or update anything else. Setting `parameter_state` to `std::nullopt`
+    // prevents the state being updated later.
+    parameter_state = std::nullopt;
     return absl::OkStatus();
   }
 
-  if (expected_timestamp != parameter_state.next_timestamp) {
-    return absl::InvalidArgumentError(absl::StrCat(
-        "Mismatching timestamps for ", parameter_name, " parameters: (",
-        parameter_state.next_timestamp, " vs ", expected_timestamp, ")"));
-  }
-
   // Update the next timestamp for the next frame.
-  parameter_state.next_timestamp = parameter_block->end_timestamp;
+  (*parameter_state)->next_timestamp = parameter_block->end_timestamp;
+  RETURN_IF_NOT_OK(CompareTimestamps(
+      expected_next_timestamp, (*parameter_state)->next_timestamp,
+      absl::StrCat("When updating states for ", parameter_name,
+                   " parameters: ")));
 
   // Clear out the parameter block, which should not be used before a new
   // one is added via `Add*ParameterBlock()`.
   parameter_block = nullptr;
 
-  // Additional update steps.
-  additional_state_updater(parameter_state);
-
   return absl::OkStatus();
 }
 
@@ -179,9 +179,9 @@ absl::Status ParametersManager::GetDownMixingParameters(
   }
   auto& demixing_state = demixing_states_iter->second;
   const auto* param_definition = demixing_state.param_definition;
-  const auto* parameter_block =
+  const auto* demixing_parameter_block =
       demixing_parameter_blocks_.at(param_definition->parameter_id_);
-  if (parameter_block == nullptr) {
+  if (demixing_parameter_block == nullptr) {
     // Failed to find a parameter block that overlaps this frame. Use the
     // default value from the parameter definition. This is OK when there are
     // no parameter blocks covering this substream. If there is only partial
@@ -196,17 +196,14 @@ absl::Status ParametersManager::GetDownMixingParameters(
     return absl::OkStatus();
   }
 
-  if (parameter_block->start_timestamp != demixing_state.next_timestamp) {
-    return absl::InvalidArgumentError(absl::StrCat(
-        "Mismatching timestamps for down-mixing parameters for "
-        "audio element ID= ",
-        audio_element_id, ": expecting", demixing_state.next_timestamp,
-        " but got ", parameter_block->start_timestamp));
-  }
+  RETURN_IF_NOT_OK(CompareTimestamps(
+      demixing_state.next_timestamp, demixing_parameter_block->start_timestamp,
+      absl::StrCat("Getting down-mixing parameters for audio element ID= ",
+                   audio_element_id, ": ")));
 
   RETURN_IF_NOT_OK(DemixingInfoParameterData::DMixPModeToDownMixingParams(
       static_cast<DemixingInfoParameterData*>(
-          parameter_block->obu->subblocks_[0].param_data.get())
+          demixing_parameter_block->obu->subblocks_[0].param_data.get())
           ->dmixp_mode,
       demixing_state.previous_w_idx, demixing_state.update_rule,
       down_mixing_params));
@@ -262,14 +259,11 @@ absl::Status ParametersManager::GetReconGainInfoParameterData(
     return absl::OkStatus();
   }
 
-  if (recon_gain_parameter_block->start_timestamp !=
-      recon_gain_state.next_timestamp) {
-    return absl::InvalidArgumentError(absl::StrCat(
-        "Mismatching timestamps for recon gain parameters for "
-        "audio element ID= ",
-        audio_element_id, ": expecting", recon_gain_state.next_timestamp,
-        " but got ", recon_gain_parameter_block->start_timestamp));
-  }
+  RETURN_IF_NOT_OK(CompareTimestamps(
+      recon_gain_state.next_timestamp,
+      recon_gain_parameter_block->start_timestamp,
+      absl::StrCat("Getting recon gain parameters for audio element ID= ",
+                   audio_element_id, ": ")));
 
   auto recon_gain_info_parameter_data_in_obu =
       static_cast<ReconGainInfoParameterData*>(
@@ -293,32 +287,33 @@ void ParametersManager::AddReconGainParameterBlock(
 
 absl::Status ParametersManager::UpdateDemixingState(
     DecodedUleb128 audio_element_id, int32_t expected_timestamp) {
-  // Additional updating steps to perform at the end of
-  // `UpdateParameterState()`.
-  auto demixing_state_updater = [](DemixingState& demixing_state) {
+  std::optional<DemixingState*> demixing_state = std::nullopt;
+  RETURN_IF_NOT_OK(UpdateParameterState(
+      audio_element_id, expected_timestamp, demixing_states_,
+      demixing_parameter_blocks_, "down-mixing", demixing_state));
+
+  // Additional updating steps to perform after `UpdateParameterState()`.
+  if (demixing_state.has_value()) {
     // Update `previous_w_idx` for the next frame.
-    demixing_state.previous_w_idx = demixing_state.w_idx;
+    (*demixing_state)->previous_w_idx = (*demixing_state)->w_idx;
 
     // Update the `update_rule` of the first frame to be "normally updating".
-    if (demixing_state.update_rule == DemixingInfoParameterData::kFirstFrame) {
-      demixing_state.update_rule = DemixingInfoParameterData::kNormal;
+    if ((*demixing_state)->update_rule ==
+        DemixingInfoParameterData::kFirstFrame) {
+      (*demixing_state)->update_rule = DemixingInfoParameterData::kNormal;
     }
-  };
-
-  return UpdateParameterState(audio_element_id, expected_timestamp,
-                              demixing_states_, demixing_parameter_blocks_,
-                              "down-mixing", demixing_state_updater);
-
+  }
   return absl::OkStatus();
 }
 
 absl::Status ParametersManager::UpdateReconGainState(
     DecodedUleb128 audio_element_id, int32_t expected_timestamp) {
-  return UpdateParameterState(
-      audio_element_id, expected_timestamp, recon_gain_states_,
-      recon_gain_parameter_blocks_, "recon gain",
-      [](auto& state) {}  // No additional updating needed.
-  );
+  std::optional<ReconGainState*> recon_gain_state = std::nullopt;
+
+  // No additional updating needed.
+  return UpdateParameterState(audio_element_id, expected_timestamp,
+                              recon_gain_states_, recon_gain_parameter_blocks_,
+                              "down-mixing", recon_gain_state);
 }
 
 }  // namespace iamf_tools

iamf/cli/parameters_manager.h

@@ -33,10 +33,13 @@ namespace iamf_tools {
  * corresponding to the same frame (with the same start/end timestamps).
  *
  * For each frame:
- *   - Parameter blocks are added via `AddDemixingParameterBlock()`,
- *   - Parameter values can be queried via `GetDownMixingParameters()`.
+ *   - Parameter blocks are added via `AddDemixingParameterBlock()` or
+ *     `AddReconGainParameterBlock()`.
+ *   - Parameter values can be queried via `GetDownMixingParameters()` or
+ *     `GetReconGainInfoParameterData()`.
  *   - Caller (usually the audio frame generator) is responsible to tell this
- *     manager to advance to the next frame via `UpdateDemixingState()`.
+ *     manager to advance to the next frame via `UpdateDemixingState()` or
+ *     `UpdateReconGainState()`.
  */
 class ParametersManager {
  public:
@@ -106,25 +109,25 @@ class ParametersManager {
    *
    * \param audio_element_id Audio Element ID whose corresponding demixing
    *        state are to be updated.
-   * \param expected_timestamp Expected timestamp of the next set of
+   * \param expected_next_timestamp Expected timestamp of the upcoming set of
    *        demixing parameter blocks.
    * \return `absl::OkStatus()` on success. A specific status on failure.
    */
   absl::Status UpdateDemixingState(DecodedUleb128 audio_element_id,
-                                   int32_t expected_timestamp);
+                                   int32_t expected_next_timestamp);
 
   /*!\brief Updates the state of recon gain parameters for an audio element.
    *
    * Also validates the timestamp is as expected.
    *
    * \param audio_element_id Audio Element ID whose corresponding recon gain
    *        state are to be updated.
-   * \param expected_timestamp Expected timestamp of the next set of
+   * \param expected_new_timestamp Expected timestamp of the upcoming set of
    *        recon gain parameter blocks.
    * \return `absl::OkStatus()` on success. A specific status on failure.
    */
   absl::Status UpdateReconGainState(DecodedUleb128 audio_element_id,
-                                    int32_t expected_timestamp);
+                                    int32_t expected_next_timestamp);
 
  private:
   // State used when generating demixing parameters for an audio element.

iamf/cli/proto_to_obu/audio_element_generator.cc

@@ -829,7 +829,7 @@ absl::Status FinalizeAmbisonicsMonoConfig(
     if (obu_substream_index ==
         AmbisonicsMonoConfig::kInactiveAmbisonicsChannelNumber) {
       LOG(INFO) << "Detected mixed-order ambisonics with  A"
-                << ambisonics_channel_number << "dropped.";
+                << ambisonics_channel_number << " dropped.";
       continue;
     }
     const DecodedUleb128 substream_id =

iamf/cli/proto_to_obu/audio_frame_generator.cc

@@ -530,10 +530,10 @@ absl::Status MaybeEncodeFramesForAudioElement(
   } while (!encoded_timestamp.has_value() && more_samples_to_encode);
 
   if (encoded_timestamp.has_value()) {
-    // An audio frame has been encoded, update the parameter manager to use
-    // the next frame of parameters.
+    // All audio frames corresponding to the audio element have been encoded;
+    // update the parameter manager to use the next frame of parameters.
     RETURN_IF_NOT_OK(parameters_manager.UpdateDemixingState(
-        audio_element_id, *encoded_timestamp));
+        audio_element_id, *encoded_timestamp + num_samples_per_frame));
   }
 
   return absl::OkStatus();

iamf/cli/tests/parameters_manager_test.cc

@@ -44,6 +44,7 @@ constexpr DecodedUleb128 kAudioElementId = 157;
 constexpr DecodedUleb128 kParameterId = 995;
 constexpr DecodedUleb128 kSecondParameterId = 996;
 constexpr DecodedUleb128 kDuration = 8;
+
 constexpr DemixingInfoParameterData::DMixPMode kDMixPMode =
     DemixingInfoParameterData::kDMixPMode3_n;
 
@@ -297,10 +298,10 @@ TEST_F(ParametersManagerTest, GetMultipleReconGainParametersSucceeds) {
   EXPECT_EQ(recon_gain_parameter_data_0.recon_gain_elements[0].recon_gain[0],
             0);
 
-  EXPECT_THAT(
-      parameters_manager_->UpdateReconGainState(kAudioElementId,
-                                                /*expected_timestamp=*/0),
-      IsOk());
+  EXPECT_THAT(parameters_manager_->UpdateReconGainState(
+                  kAudioElementId,
+                  /*expected_next_timestamp=*/kDuration),
+              IsOk());
 
   // Second recon gain parameter block.
   ASSERT_THAT(AddOneReconGainParameterBlock(
@@ -326,7 +327,7 @@ TEST_F(ParametersManagerTest, GetMultipleReconGainParametersSucceeds) {
   // offset by the duration of the parameter block.
   EXPECT_THAT(parameters_manager_->UpdateReconGainState(
                   kAudioElementId,
-                  /*expected_timestamp=*/kDuration),
+                  /*expected_next_timestamp=*/kDuration + kDuration),
               IsOk());
 }
 
@@ -384,10 +385,10 @@ TEST_F(ParametersManagerTest, ParameterBlocksRunOutReturnsDefault) {
                                                            down_mixing_params),
               IsOk());
 
-  EXPECT_THAT(
-      parameters_manager_->UpdateDemixingState(kAudioElementId,
-                                               /*expected_timestamp=*/0),
-      IsOk());
+  EXPECT_THAT(parameters_manager_->UpdateDemixingState(
+                  kAudioElementId,
+                  /*expected_next_timestamp=*/kDuration),
+              IsOk());
 
   // Get the parameters for the second time. Since there is only one
   // parameter block and is already used up the previous time, the function
@@ -406,12 +407,13 @@ TEST_F(ParametersManagerTest, ParameterBlocksRunOutReturnsDefault) {
   EXPECT_EQ(down_mixing_params.w_idx_used, 10);
   EXPECT_FLOAT_EQ(down_mixing_params.w, 0.5);
 
-  // `UpdateDemixingState()` also succeeds, because technically there's
-  // nothing to update.
-  EXPECT_THAT(
-      parameters_manager_->UpdateDemixingState(kAudioElementId,
-                                               /*expected_timestanmp=*/8),
-      IsOk());
+  // `UpdateDemixingState()` also succeeds with some arbitrary timestamp,
+  // because technically there's nothing to update.
+  const DecodedUleb128 kArbitraryTimestamp = 972;
+  EXPECT_THAT(parameters_manager_->UpdateDemixingState(
+                  kAudioElementId,
+                  /*expected_next_timestamp=*/kArbitraryTimestamp),
+              IsOk());
 }
 
 TEST_F(ParametersManagerTest, ParameterIdNotFoundReturnsDefault) {
@@ -465,10 +467,10 @@ TEST_F(ParametersManagerTest, GetDownMixingParametersTwiceDifferentW) {
   ASSERT_THAT(parameters_manager_->GetDownMixingParameters(kAudioElementId,
                                                            down_mixing_params),
               IsOk());
-  EXPECT_THAT(
-      parameters_manager_->UpdateDemixingState(kAudioElementId,
-                                               /*expected_timestamp=*/0),
-      IsOk());
+  EXPECT_THAT(parameters_manager_->UpdateDemixingState(
+                  kAudioElementId,
+                  /*expected_next_timestamp=*/kDuration),
+              IsOk());
 
   // The first time `w_idx` is 0, and the corresponding `w` is 0.
   const double kWFirst = 0.0;
@@ -572,10 +574,10 @@ TEST_F(ParametersManagerTest,
   ASSERT_THAT(parameters_manager_->GetDownMixingParameters(kAudioElementId,
                                                            down_mixing_params),
               IsOk());
-  EXPECT_THAT(
-      parameters_manager_->UpdateDemixingState(kAudioElementId,
-                                               /*expected_timestamp=*/0),
-      IsOk());
+  EXPECT_THAT(parameters_manager_->UpdateDemixingState(
+                  kAudioElementId,
+                  /*expected_next_timestamp=*/kDuration),
+              IsOk());
   EXPECT_FLOAT_EQ(down_mixing_params.w, kWFirst);
 
   // Add the parameter block for the first audio element corresponding to the
@@ -614,7 +616,9 @@ TEST_F(ParametersManagerTest, DemixingParamDefinitionIsNotAvailableForWrongId) {
               IsOk());
 
   // `UpdateDemixingState()` also succeeds.
-  EXPECT_THAT(parameters_manager_->UpdateDemixingState(kWrongAudioElementId, 0),
+  EXPECT_THAT(parameters_manager_->UpdateDemixingState(
+                  kWrongAudioElementId,
+                  /*expected_next_timestamp=*/kDuration),
               IsOk());
 }
 
@@ -624,12 +628,12 @@ TEST_F(ParametersManagerTest, UpdateFailsWithWrongTimestamps) {
   parameters_manager_->AddDemixingParameterBlock(
       &demixing_parameter_blocks_[0]);
 
-  // The first frame starts with timestamp = 0, so updating with a different
+  // The second frame starts with timestamp = 8, so updating with a different
   // timestamp fails.
-  const int32_t kWrongTimestamp = 8;
-  EXPECT_FALSE(
-      parameters_manager_->UpdateDemixingState(kAudioElementId, kWrongTimestamp)
-          .ok());
+  const int32_t kWrongNextTimestamp = 17;
+  EXPECT_FALSE(parameters_manager_
+                   ->UpdateDemixingState(kAudioElementId, kWrongNextTimestamp)
+                   .ok());
 }
 
 TEST_F(ParametersManagerTest, UpdateNotValidatingWhenParameterIdNotFound) {