From 678264437b852d1b603625dfd4ce2b7e0b8211fc Mon Sep 17 00:00:00 2001 From: Filyus Date: Sun, 4 Oct 2026 16:14:02 +0500 Subject: [PATCH] Make KD-tree decoder memory linear in the point dimension DynamicIntegerPointsKdTreeDecoder allocates memory quadratic in the point dimension, up front, in its constructor. The dimension is the total number of components over all attributes, and a stream sets it at five bytes per attribute. This change makes the decoder's memory linear in the dimension. The decoded output is unchanged. Decoder memory by dimension, main always allocated against this change at most: 3 components: 2.3 KB -> 1.6 KB 64 components: 1 MB -> 33 KB 255 components: 17 MB -> 133 KB 2,048 components: 1.1 GB -> 1.1 MB 6,375 components: 10 GB -> 3.3 MB A 143-byte point cloud with 25 attributes of 255 components makes main commit about 10 GB before it reads any KD-tree data, and then fail. With this change it fails without that allocation. A valid 8,372-byte stream of three points with 2,048 components decodes with about 5 MB in 80 ms, against 1.1 GB and 1 to 1.9 s on main. Ordinary point clouds decode as fast or faster, from 1% faster with 3 components to 30% faster with 64. The decoder kept a full copy of the node's base and levels for every level of the tree, which is 32 * dimension levels deep. A split changes only one base value and one level. The decoder now keeps a single base and levels vector, records the old values at each split, and restores them when the walk returns to a shallower node. --- .../dynamic_integer_points_kd_tree_decoder.h | 92 ++++++++++++++----- .../point_cloud_kd_tree_encoding_test.cc | 23 +++++ 2 files changed, 94 insertions(+), 21 deletions(-) diff --git a/src/draco/compression/point_cloud/algorithms/dynamic_integer_points_kd_tree_decoder.h b/src/draco/compression/point_cloud/algorithms/dynamic_integer_points_kd_tree_decoder.h index 55bafe7c4..b5b8dd647 100644 --- a/src/draco/compression/point_cloud/algorithms/dynamic_integer_points_kd_tree_decoder.h +++ b/src/draco/compression/point_cloud/algorithms/dynamic_integer_points_kd_tree_decoder.h @@ -88,11 +88,7 @@ class DynamicIntegerPointsKdTreeDecoder { num_decoded_points_(0), dimension_(dimension), p_(dimension, 0), - axes_(dimension, 0), - // Init the stack with the maximum depth of the tree. - // +1 for a second leaf. - base_stack_(32 * dimension + 1, VectorUint32(dimension, 0)), - levels_stack_(32 * dimension + 1, VectorUint32(dimension, 0)) {} + axes_(dimension, 0) {} // Decodes an integer point cloud from |buffer|. Optional |oit_max_points| can // be used to tell the decoder the maximum number of points accepted by the @@ -130,14 +126,25 @@ class DynamicIntegerPointsKdTreeDecoder { struct DecodingStatus { DecodingStatus(uint32_t num_remaining_points_, uint32_t last_axis_, - uint32_t stack_pos_) + uint32_t depth_) : num_remaining_points(num_remaining_points_), last_axis(last_axis_), - stack_pos(stack_pos_) {} + depth(depth_) {} uint32_t num_remaining_points; uint32_t last_axis; - uint32_t stack_pos; // used to get base and levels + uint32_t depth; // used to get base and levels + }; + + // What one split changed in |base_| and |levels_|, kept so the walk can put + // it back when it returns to a shallower node. The node that split keeps + // the new level of |axis|. The node one level deeper also has the split + // bit set in its base. + struct SplitUndo { + uint32_t depth; + uint32_t axis; + uint32_t old_level; + uint32_t old_base; }; uint32_t bit_length_; @@ -150,8 +157,11 @@ class DynamicIntegerPointsKdTreeDecoder { HalfDecoder half_decoder_; VectorUint32 p_; VectorUint32 axes_; - std::vector base_stack_; - std::vector levels_stack_; + // The base and levels of the node being decoded, and the splits on the + // walk's current path, deepest last. + VectorUint32 base_; + VectorUint32 levels_; + std::vector undo_; }; // Decodes a point cloud from |buffer|. @@ -252,8 +262,22 @@ template bool DynamicIntegerPointsKdTreeDecoder::DecodeInternal( uint32_t num_points, OutputIteratorT &oit) { typedef DecodingStatus Status; - base_stack_[0] = VectorUint32(dimension_, 0); - levels_stack_[0] = VectorUint32(dimension_, 0); + // The tree can be 32 * dimension levels deep, so keeping a base and levels + // vector per level takes memory quadratic in the dimension. The walk keeps + // one of each instead. A split at depth d raises one level of the node at + // d, and the node at d + 1 has the same levels and one more base bit set. + // Each split is logged in |undo_|, and on reaching a node at depth t the + // walk undoes the splits deeper than t and the base bit of the last split + // at t. + // + // A split takes its node off the top of |status_stack|, where every node + // below is shallower, and pushes its halves at d and d + 1. So the depths + // on the stack strictly increase towards its top, and a node is reached + // only after everything deeper than it has been decoded. + base_.assign(dimension_, 0); + levels_.assign(dimension_, 0); + undo_.clear(); + uint32_t depth = 0; DecodingStatus init_status(num_points, 0, 0); std::stack status_stack; status_stack.push(init_status); @@ -265,9 +289,29 @@ bool DynamicIntegerPointsKdTreeDecoder::DecodeInternal( const uint32_t num_remaining_points = status.num_remaining_points; const uint32_t last_axis = status.last_axis; - const uint32_t stack_pos = status.stack_pos; - const VectorUint32 &old_base = base_stack_[stack_pos]; - const VectorUint32 &levels = levels_stack_[stack_pos]; + // A node at the current depth is the second half of the split just made, + // and nothing needs to be undone for it. + if (status.depth != depth) { + if (status.depth > depth) { + return false; + } + while (!undo_.empty()) { + const SplitUndo &split = undo_.back(); + if (split.depth < status.depth) { + break; + } + base_[split.axis] = split.old_base; + if (split.depth == status.depth) { + // The node's own split, whose level stays raised. + break; + } + levels_[split.axis] = split.old_level; + undo_.pop_back(); + } + depth = status.depth; + } + const VectorUint32 &old_base = base_; + const VectorUint32 &levels = levels_; if (num_remaining_points > num_points) { return false; @@ -325,8 +369,6 @@ bool DynamicIntegerPointsKdTreeDecoder::DecodeInternal( const int num_remaining_bits = bit_length_ - level; const uint32_t modifier = 1 << (num_remaining_bits - 1); - base_stack_[stack_pos + 1] = old_base; // copy - base_stack_[stack_pos + 1][axis] += modifier; // new base const int incoming_bits = MostSignificantBit(num_remaining_points); @@ -347,14 +389,22 @@ bool DynamicIntegerPointsKdTreeDecoder::DecodeInternal( } } - levels_stack_[stack_pos][axis] += 1; - levels_stack_[stack_pos + 1] = levels_stack_[stack_pos]; // copy + if (depth == std::numeric_limits::max()) { + return false; + } + // Both halves see |axis| one level deeper. The first half stays at this + // depth and keeps the base, the second goes one deeper with the split bit + // set in its base. + undo_.push_back({depth, axis, levels_[axis], base_[axis]}); + levels_[axis] += 1; + base_[axis] += modifier; if (first_half) { - status_stack.push(DecodingStatus(first_half, axis, stack_pos)); + status_stack.push(DecodingStatus(first_half, axis, depth)); } if (second_half) { - status_stack.push(DecodingStatus(second_half, axis, stack_pos + 1)); + status_stack.push(DecodingStatus(second_half, axis, depth + 1)); } + ++depth; } return true; } diff --git a/src/draco/compression/point_cloud/point_cloud_kd_tree_encoding_test.cc b/src/draco/compression/point_cloud/point_cloud_kd_tree_encoding_test.cc index 1ca02d58e..7f59e8124 100644 --- a/src/draco/compression/point_cloud/point_cloud_kd_tree_encoding_test.cc +++ b/src/draco/compression/point_cloud/point_cloud_kd_tree_encoding_test.cc @@ -140,6 +140,29 @@ TEST_F(PointCloudKdTreeEncodingTest, TestRejectMismatchedLegacyPointCount) { EXPECT_FALSE(decoder.DecodePointCloudFromBuffer(&buffer).ok()); } +TEST_F(PointCloudKdTreeEncodingTest, TestRejectTruncatedHighDimensionalStream) { + // A KD-tree point cloud with one point and 25 float attributes of 255 + // components each, for a total dimension of 6375. The stream ends after the + // compression level, before any of the KD-tree data. The decoder must fail + // without allocating stacks for a walk it never reads. + std::vector data = {'D', 'R', 'A', 'C', 'O', 2, 3, 0, 1, 0, 0, + // One point, one attributes decoder and 25 + // attributes. + 1, 0, 0, 0, 1, 25}; + for (uint8_t i = 0; i < 25; ++i) { + // Generic, float32, 255 components, not normalized, unique id |i|. + const uint8_t attribute[] = {4, 9, 255, 0, i}; + data.insert(data.end(), attribute, attribute + sizeof(attribute)); + } + // Compression level 0. + data.push_back(0); + + DecoderBuffer buffer; + buffer.Init(reinterpret_cast(data.data()), data.size()); + Decoder decoder; + EXPECT_FALSE(decoder.DecodePointCloudFromBuffer(&buffer).ok()); +} + // test higher dimensions with more attributes TEST_F(PointCloudKdTreeEncodingTest, TestIntKdTreeEncodingHigherDimension) { constexpr int num_points = 120;