Fix push_array_to_array stdlib method

cpp/src/barretenberg/stdlib/primitives/field/array.hpp

@@ -132,36 +132,44 @@ template <typename Composer, size_t size_1, size_t size_2>
 void push_array_to_array(std::array<field_t<Composer>, size_1> const& source,
                          std::array<field_t<Composer>, size_2>& target)
 {
-    // TODO: inefficient to get length this way within this function. Probably best to inline the checks that we need
-    // into the below loops directly.
+    ASSERT(target.size() >= source.size());
+    field_t<Composer> source_length = array_length<Composer>(source);
     field_t<Composer> target_length = array_length<Composer>(target);
-    field_t<Composer> target_capacity = field_t<Composer>(target.size());
-    const field_t<Composer> overflow_capacity = target_capacity + 1;
 
-    // ASSERT(uint256_t(target_capacity.get_value()) + 1 >
-    //        uint256_t(target_length.get_value()) + uint256_t(source_length.get_value()));
-
-    field_t<Composer> j_ct = 0; // circuit-type index for the inner loop
-    field_t<Composer> next_target_index = target_length;
-    for (size_t i = 0; i < source.size(); ++i) {
-        auto& s = source[i];
-
-        // Triangular loop:
-        for (size_t j = i; j < target.size() - source.size() + i + 1; ++j) {
-            auto& t = target[j];
-
-            bool_t<Composer> at_next_index = j_ct == next_target_index;
-
-            t = field_t<Composer>::conditional_assign(at_next_index, s, t);
+    // Check if the `source` array is too large vs the remaining capacity of the `target` array
+    auto source_length_safe = safe_uint_t<Composer>(source_length, sizeof(size_t) * 8, "source_array_len");
+    auto target_length_safe = safe_uint_t<Composer>(target_length, sizeof(size_t) * 8, "target_array_len");
+    auto target_capacity_safe =
+        safe_uint_t<Composer>(field_t<Composer>(target.size()), sizeof(size_t) * 8, "target_array_capacity");
+    target_capacity_safe.subtract((source_length_safe + target_length_safe),
+                                  sizeof(size_t) * 8,
+                                  "push_array_to_array target array capacity exceeded");
+
+    // Ensure that all the elements after the first zero-element in target are zero
+    field_t<Composer> num_non_zero_in_target = 0;
+    for (size_t i = 0; i < size_2; i++) {
+        num_non_zero_in_target += !target[i].is_zero();
+    }
+    num_non_zero_in_target.assert_equal(target_length_safe, "non-zero element in target array after a zero element");
 
-            j_ct++;
+    // Ensure that all the elements after the first zero-element in source are zero
+    field_t<Composer> num_non_zero_in_source = 0;
+    for (size_t i = 0; i < size_1; i++) {
+        num_non_zero_in_source += !source[i].is_zero();
+    }
+    num_non_zero_in_source.assert_equal(source_length_safe, "non-zero element in source array after a zero element");
+
+    // Finally, insert only non-zero elements from source into target
+    bool_t<Composer> composer_error_found = !source_length.get_context()->err().empty();
+    for (size_t i = 0; i < source.size(); i++) {
+        bool_t<Composer> found = false;
+
+        for (size_t j = i; j < target.size(); j++) {
+            bool_t<Composer> is_target_non_zero = !target[j].is_zero();
+            target[j] = field_t<Composer>::conditional_assign(
+                is_target_non_zero, target[j], source[i] * !found * !composer_error_found);
+            found |= !is_target_non_zero;
         }
-
-        next_target_index++;
-
-        next_target_index.assert_not_equal(overflow_capacity, "push_array_to_array target array capacity exceeded");
-
-        j_ct = i + 1;
     }
 }
 

cpp/src/barretenberg/stdlib/primitives/field/field.test.cpp

@@ -9,6 +9,7 @@
 #include "barretenberg/plonk/composer/ultra_composer.hpp"
 #include "barretenberg/plonk/composer/turbo_composer.hpp"
 #include "barretenberg/numeric/random/engine.hpp"
+#include "barretenberg/common/streams.hpp"
 
 using namespace bonk;
 
@@ -1180,61 +1181,136 @@ template <typename Composer> class stdlib_field : public testing::Test {
         EXPECT_EQ(proof_result, true);
     };
 
-    static void test_push_array_to_array()
+    template <size_t size_1, size_t size_2>
+    static auto test_push_array_to_array_helper(Composer& composer,
+                                                std::array<fr, size_1> const& source,
+                                                std::array<fr, size_2> const& target,
+                                                std::array<fr, size_2> const& expected_target)
     {
-        Composer composer = Composer();
+        std::array<field_ct, size_1> source_ct;
+        std::array<field_ct, size_2> target_ct;
+        for (size_t i = 0; i < source.size(); i++) {
+            source_ct[i] = witness_ct(&composer, source[i]);
+        }
+        for (size_t i = 0; i < target.size(); i++) {
+            target_ct[i] = witness_ct(&composer, target[i]);
+        }
 
-        std::array<field_ct, 4> source = { 1, 2 };
-        std::array<field_ct, 8> target = { 3, 4, 6, 8, 0, 0, 0, 0 };
-        push_array_to_array<Composer>(source, target);
+        push_array_to_array<Composer>(source_ct, target_ct);
 
         // Check that the source array has been inserted into the first available index of the target array.
-        ASSERT(target[0].get_value() == 3);
-        ASSERT(target[1].get_value() == 4);
-        ASSERT(target[2].get_value() == 6);
-        ASSERT(target[3].get_value() == 8);
-        ASSERT(target[4].get_value() == 1);
-        ASSERT(target[5].get_value() == 2);
-        ASSERT(target[6].get_value() == 0);
-        ASSERT(target[7].get_value() == 0);
+        for (size_t i = 0; i < target.size(); i++) {
+            ASSERT(target_ct[i].get_value() == expected_target[i]);
+        }
 
-        auto prover = composer.create_prover();
-        auto verifier = composer.create_verifier();
-        auto proof = prover.construct_proof();
-        info("composer gates = ", composer.get_num_gates());
-        bool proof_result = verifier.verify_proof(proof);
-        EXPECT_EQ(proof_result, true);
+        bool proof_result = false;
+        if (composer.err().empty()) {
+            auto prover = composer.create_prover();
+            auto verifier = composer.create_verifier();
+            auto proof = prover.construct_proof();
+            info("composer gates = ", composer.get_num_gates());
+            proof_result = verifier.verify_proof(proof);
+        }
+
+        return std::make_pair(proof_result, composer.err());
     }
 
-    static void test_push_array_to_array_full()
+    static void test_push_array_to_array()
     {
-        Composer composer = Composer();
+        {
+            Composer composer = Composer();
 
-        std::array<fr, 4> source = { 1, 2 };
-        std::array<fr, 8> target = { 3, 4, 6, 8, 7, 9, 5, 0 };
-        std::array<field_ct, 4> source_ct;
-        std::array<field_ct, 8> target_ct;
-        for (size_t i = 0; i < source.size(); i++) {
-            source_ct[i] = witness_ct(&composer, source[i]);
+            std::array<fr, 4> source = { 1, 0, 0, 0 };
+            std::array<fr, 4> target = { 3, 0, 0, 0 };
+            std::array<fr, 4> expected_target = { 3, 1, 0, 0 };
+            bool proof_res;
+            std::string error;
+            std::tie(proof_res, error) =
+                test_push_array_to_array_helper<4, 4>(composer, source, target, expected_target);
+
+            EXPECT_EQ(proof_res, true);
         }
-        for (size_t i = 0; i < target.size(); i++) {
-            target_ct[i] = witness_ct(&composer, target[i]);
+        {
+            Composer composer = Composer();
+
+            std::array<fr, 4> source = { 1, 2, 3, 0 };
+            std::array<fr, 4> target = { 0, 0, 0, 0 };
+            std::array<fr, 4> expected_target = { 1, 2, 3, 0 };
+            bool proof_res;
+            std::string error;
+            std::tie(proof_res, error) =
+                test_push_array_to_array_helper<4, 4>(composer, source, target, expected_target);
+
+            EXPECT_EQ(proof_res, true);
         }
+        {
+            Composer composer = Composer();
 
-        push_array_to_array<Composer>(source_ct, target_ct);
+            std::array<fr, 3> source = { 1, 2, 3 };
+            std::array<fr, 6> target = { 4, 5, 6, 0, 0, 0 };
+            std::array<fr, 6> expected_target = { 4, 5, 6, 1, 2, 3 };
+            bool proof_res;
+            std::string error;
+            std::tie(proof_res, error) =
+                test_push_array_to_array_helper<3, 6>(composer, source, target, expected_target);
 
-        // Check that the target array is unchanged as the push operation failed
-        ASSERT(target_ct[0].get_value() == 3);
-        ASSERT(target_ct[1].get_value() == 4);
-        ASSERT(target_ct[2].get_value() == 6);
-        ASSERT(target_ct[3].get_value() == 8);
-        ASSERT(target_ct[4].get_value() == 7);
-        ASSERT(target_ct[5].get_value() == 9);
-        ASSERT(target_ct[6].get_value() == 5);
-        ASSERT(target_ct[7].get_value() == 0);
+            EXPECT_EQ(proof_res, true);
+        }
+        {
+            Composer composer = Composer();
 
-        EXPECT_EQ(composer.failed(), true);
-        EXPECT_EQ(composer.err(), "push_array_to_array target array capacity exceeded");
+            std::array<fr, 5> source = { 1, 2, 3, 4, 5 };
+            std::array<fr, 5> target = { 5, 6, 7, 8, 9 };
+            std::array<fr, 5> expected_target = { 5, 6, 7, 8, 9 };
+            bool proof_res;
+            std::string error;
+            std::tie(proof_res, error) = test_push_array_to_array_helper(composer, source, target, expected_target);
+
+            EXPECT_EQ(proof_res, false);
+            EXPECT_EQ(
+                error,
+                "safe_uint_t range constraint failure: subtract: push_array_to_array target array capacity exceeded");
+        }
+        {
+            Composer composer = Composer();
+
+            std::array<fr, 4> source = { 1, 0, 2, 3 };
+            std::array<fr, 6> target = { 4, 5, 6, 7, 8, 0 };
+            std::array<fr, 6> expected_target = { 4, 5, 6, 7, 8, 0 };
+            bool proof_res;
+            std::string error;
+            std::tie(proof_res, error) = test_push_array_to_array_helper(composer, source, target, expected_target);
+
+            EXPECT_EQ(proof_res, false);
+            EXPECT_EQ(error, "non-zero element in source array after a zero element");
+        }
+        {
+            Composer composer = Composer();
+
+            std::array<fr, 3> source = { 1, 2, 3 };
+            std::array<fr, 6> target = { 4, 5, 0, 6, 7, 8 };
+            std::array<fr, 6> expected_target = { 4, 5, 0, 6, 7, 8 };
+            bool proof_res;
+            std::string error;
+            std::tie(proof_res, error) = test_push_array_to_array_helper(composer, source, target, expected_target);
+
+            EXPECT_EQ(proof_res, false);
+            EXPECT_EQ(error, "non-zero element in target array after a zero element");
+        }
+        {
+            // captures the last composer error
+            Composer composer = Composer();
+
+            std::array<fr, 3> source = { 1, 0, 3 };
+            std::array<fr, 6> target = { 4, 5, 0, 6, 7, 8 };
+            std::array<fr, 6> expected_target = { 4, 5, 0, 6, 7, 8 };
+            bool proof_res;
+            std::string error;
+            std::tie(proof_res, error) = test_push_array_to_array_helper(composer, source, target, expected_target);
+
+            EXPECT_EQ(proof_res, false);
+            EXPECT_EQ(error, "non-zero element in target array after a zero element");
+        }
     }
 
     class MockClass {
@@ -1470,8 +1546,4 @@ TYPED_TEST(stdlib_field, test_array_push_array_to_array)
 {
     TestFixture::test_push_array_to_array();
 }
-TYPED_TEST(stdlib_field, test_array_push_array_to_array_full)
-{
-    TestFixture::test_push_array_to_array_full();
-}
 } // namespace test_stdlib_field