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;