diff --git a/NEWS b/NEWS index 4ebf50fe6a..a2d37fc588 100644 --- a/NEWS +++ b/NEWS @@ -96,6 +96,14 @@ New Functionality ``disable_analyzer()`` at a later point when it finds the condition due to which it vetoed fulfilled (which may be never). +- Add support for iterating over indices and values of a vector using the + same syntax as used for iterating over key-value pairs of tables, where + ``value`` will be set to ``vec[idx]``. + + local vec = vector("zero", "one", "two"); + for ( idx, value in vec ) + print idx, value; + Changed Functionality --------------------- diff --git a/src/Stmt.cc b/src/Stmt.cc index 0d5a47afaf..7d2b8d93c9 100644 --- a/src/Stmt.cc +++ b/src/Stmt.cc @@ -1290,23 +1290,29 @@ ForStmt::ForStmt(IDPList* arg_loop_vars, ExprPtr loop_expr, IDPtr val_var) { value_var = std::move(val_var); - if ( e->GetType()->IsTable() ) - { - const auto& yield_type = e->GetType()->AsTableType()->Yield(); + auto t = e->GetType(); + zeek::TypePtr yield_type; - // Verify value_vars type if its already been defined - if ( value_var->GetType() ) - { - if ( ! same_type(value_var->GetType(), yield_type) ) - value_var->GetType()->Error("type clash in iteration", yield_type.get()); - } - else - { - add_local(value_var, yield_type, INIT_NONE, nullptr, nullptr, VAR_REGULAR); - } + if ( t->IsTable() ) + yield_type = t->AsTableType()->Yield(); + + else if ( t->Tag() == TYPE_VECTOR ) + yield_type = t->AsVectorType()->Yield(); + + else + { + e->Error("key value for loops only support iteration over tables or vectors"); + return; + } + + // Verify value_vars type if it's already been defined + if ( value_var->GetType() ) + { + if ( ! same_type(value_var->GetType(), yield_type) ) + value_var->GetType()->Error("type clash in iteration", yield_type.get()); } else - e->Error("key value for loops only support iteration over tables"); + add_local(value_var, yield_type, INIT_NONE, nullptr, nullptr, VAR_REGULAR); } ForStmt::~ForStmt() @@ -1358,8 +1364,11 @@ ValPtr ForStmt::DoExec(Frame* f, Val* v, StmtFlowType& flow) if ( ! raw_vv[i] ) continue; - // Set the loop variable to the current index, and make - // another pass over the loop body. + // Set the loop variable to the current index, the value variable + // to the current value, and make another pass over the loop body. + if ( value_var ) + f->SetElement(value_var, vv->ValAt(i)); + f->SetElement((*loop_vars)[0], val_mgr->Count(i)); flow = FLOW_NEXT; ret = body->Exec(f, flow); diff --git a/testing/btest/Baseline/language.vector-for-value/out b/testing/btest/Baseline/language.vector-for-value/out new file mode 100644 index 0000000000..324c56fc9d --- /dev/null +++ b/testing/btest/Baseline/language.vector-for-value/out @@ -0,0 +1,13 @@ +### BTest baseline data generated by btest-diff. Do not edit. Use "btest -U/-u" to update. Requires BTest >= 0.63. +0, zero +1, one +2, two +4, four +0, zero +1, one +2, two +4, four +zero +one +two +four diff --git a/testing/btest/language/vector-for-value.zeek b/testing/btest/language/vector-for-value.zeek new file mode 100644 index 0000000000..c139bf5a3e --- /dev/null +++ b/testing/btest/language/vector-for-value.zeek @@ -0,0 +1,20 @@ +# @TEST-EXEC: zeek -b %INPUT >out +# @TEST-EXEC: btest-diff out + +local vec: vector of string = { "zero", "one", "two" }; +vec[4] = "four"; + +for ( i, v in vec ) + { + print i, v; + } + +for ( [i], v in vec ) + { + print i, v; + } + +for ( _, v in vec ) + { + print v; + }