diff --git a/src/Expr.cc b/src/Expr.cc index 768d70ef0b..ea8253e347 100644 --- a/src/Expr.cc +++ b/src/Expr.cc @@ -3201,24 +3201,26 @@ void IndexExpr::Assign(Frame* f, Val* v, Opcode op) case TYPE_VECTOR: { const ListVal *lv = v2->AsListVal(); + VectorVal* v1_vect = v1->AsVectorVal(); if ( lv->Length() > 1 ) { - int len = v1->AsVectorVal()->Size(); + int len = v1_vect->Size(); bro_int_t first = get_slice_index(lv->Index(0)->CoerceToInt(), len); bro_int_t last = get_slice_index(lv->Index(1)->CoerceToInt(), len); - int slice_length = last - first; - const VectorVal *v_vect = v->AsVectorVal(); - if ( slice_length != v_vect->Size()) - RuntimeError("vector being assigned to slice does not match size of slice"); - else if ( slice_length >= 0 ) + // Remove the elements from the vector within the slice + for ( int idx = first; idx < last; idx++ ) + v1_vect->Remove(first); + + // Insert the new elements starting at the first position + VectorVal* v_vect = v->AsVectorVal(); + for ( int idx = 0; idx < v_vect->Size(); idx++, first++ ) { - for ( int idx = first; idx < last; idx++ ) - v1->AsVectorVal()->Assign(idx, v_vect->Lookup(idx - first)->Ref(), op); + v1_vect->Insert(first, v_vect->Lookup(idx)->Ref()); } } - else if ( !v1->AsVectorVal()->Assign(v2, v, op)) + else if ( !v1_vect->Assign(v2, v, op) ) { if ( v ) { diff --git a/src/Val.cc b/src/Val.cc index bb9a3d1601..0a6c7dd985 100644 --- a/src/Val.cc +++ b/src/Val.cc @@ -3407,6 +3407,44 @@ bool VectorVal::AssignRepeat(unsigned int index, unsigned int how_many, return true; } +bool VectorVal::Insert(unsigned int index, Val* element) + { + if ( element && + ! same_type(element->Type(), vector_type->YieldType(), 0) ) + { + Unref(element); + return false; + } + + vector::iterator it; + if ( index < val.vector_val->size() ) + it = std::next(val.vector_val->begin(), index); + else + it = val.vector_val->end(); + + // Note: we do *not* Ref() the element, if any, at this point. + // AssignExpr::Eval() already does this; other callers must remember + // to do it similarly. + val.vector_val->insert(it, element); + + Modified(); + return true; + } + +bool VectorVal::Remove(unsigned int index) + { + if ( index >= val.vector_val->size() ) + return false; + + Val* val_at_index = (*val.vector_val)[index]; + auto it = std::next(val.vector_val->begin(), index); + val.vector_val->erase(it); + Unref(val_at_index); + + Modified(); + return true; + } + int VectorVal::AddTo(Val* val, int /* is_first_init */) const { if ( val->Type()->Tag() != TYPE_VECTOR ) diff --git a/src/Val.h b/src/Val.h index 2890c4c5e8..20302e3a9d 100644 --- a/src/Val.h +++ b/src/Val.h @@ -1194,6 +1194,12 @@ public: // Won't shrink size. unsigned int ResizeAtLeast(unsigned int new_num_elements); + // Insert an element at a specific position into the underlying vector. + bool Insert(unsigned int index, Val* element); + + // Removes an element or a range of elements from a specific position. + bool Remove(unsigned int start_index); + protected: friend class Val; VectorVal() { } diff --git a/testing/btest/Baseline/language.vector/out b/testing/btest/Baseline/language.vector/out index cfd9b75413..2955eda26c 100644 --- a/testing/btest/Baseline/language.vector/out +++ b/testing/btest/Baseline/language.vector/out @@ -65,3 +65,5 @@ slicing (PASS) slicing (PASS) slicing assignment (PASS) slicing assignment (PASS) +slicing assignment grow (PASS) +slicing assignment shrink (PASS) diff --git a/testing/btest/language/vector.zeek b/testing/btest/language/vector.zeek index e8b3f68e76..ea330f3842 100644 --- a/testing/btest/language/vector.zeek +++ b/testing/btest/language/vector.zeek @@ -179,4 +179,8 @@ event zeek_init() test_case( "slicing assignment", all_set(v17 == vector(6, 2, 3, 4, 5)) ); v17[2:4] = vector(7, 8); test_case( "slicing assignment", all_set(v17 == vector(6, 2, 7, 8, 5)) ); + v17[2:4] = vector(9, 10, 11); + test_case( "slicing assignment grow", all_set(v17 == vector(6, 2, 9, 10, 11, 5)) ); + v17[2:5] = vector(9); + test_case( "slicing assignment shrink", all_set(v17 == vector(6, 2, 9, 5)) ); }