GH-393: Add slice notation for vectors

This commit is contained in:
Tim Wojtulewicz 2019-06-12 13:48:56 -07:00
parent 0ab1f0fe25
commit 0af79a7a16
4 changed files with 44 additions and 8 deletions

View file

@ -2913,8 +2913,8 @@ IndexExpr::IndexExpr(Expr* arg_op1, ListExpr* arg_op2, bool is_slice)
if ( is_slice )
{
if ( ! IsString(op1->Type()->Tag()) )
ExprError("slice notation indexing only supported for strings currently");
if ( ! IsString(op1->Type()->Tag()) && ! IsVector(op1->Type()->Tag()) )
ExprError("slice notation indexing only supported for strings and vectors currently");
}
else if ( IsString(op1->Type()->Tag()) )
@ -2937,8 +2937,7 @@ IndexExpr::IndexExpr(Expr* arg_op1, ListExpr* arg_op2, bool is_slice)
else if ( ! op1->Type()->YieldType() )
{
if ( IsString(op1->Type()->Tag()) &&
match_type == MATCHES_INDEX_SCALAR )
if ( IsString(op1->Type()->Tag()) && match_type == MATCHES_INDEX_SCALAR )
SetType(base_type(TYPE_STRING));
else
// It's a set - so indexing it yields void. We don't
@ -3104,7 +3103,33 @@ Val* IndexExpr::Fold(Val* v1, Val* v2) const
switch ( v1->Type()->Tag() ) {
case TYPE_VECTOR:
v = v1->AsVectorVal()->Lookup(v2);
{
VectorVal* vect = v1->AsVectorVal();
const ListVal* lv = v2->AsListVal();
if ( lv->Length() == 1 )
v = vect->Lookup(v2);
else
{
int len = vect->Size();
VectorVal* result = nullptr;
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 sub_length = last - first;
if ( sub_length >= 0 )
{
result = new VectorVal(vect->Type()->AsVectorType());
result->Resize(sub_length);
for ( int idx = first; idx < last; idx++ )
result->Assign(idx - first, vect->Lookup(idx)->Ref());
}
return result;
}
}
break;
case TYPE_TABLE:

View file

@ -1773,10 +1773,12 @@ int VectorType::MatchesIndex(ListExpr*& index) const
{
expr_list& el = index->Exprs();
if ( el.length() != 1 )
if ( el.length() != 1 && el.length() != 2)
return DOES_NOT_MATCH_INDEX;
if ( el[0]->Type()->Tag() == TYPE_VECTOR )
if ( el.length() == 2 )
return MATCHES_INDEX_VECTOR;
else if ( el[0]->Type()->Tag() == TYPE_VECTOR )
return (IsIntegral(el[0]->Type()->YieldType()->Tag()) ||
IsBool(el[0]->Type()->YieldType()->Tag())) ?
MATCHES_INDEX_VECTOR : DOES_NOT_MATCH_INDEX;

View file

@ -58,3 +58,7 @@ access element (PASS)
&& operator (PASS)
|| operator (PASS)
+= operator (PASS)
slicing (PASS)
slicing (PASS)
slicing (PASS)
slicing (PASS)

View file

@ -168,5 +168,10 @@ event zeek_init()
v16 += 40;
test_case( "+= operator", all_set(v16 == vector( 10, 20, 30, 40 )) );
# Slicing tests.
local v17 = vector( 1, 2, 3, 4, 5 );
test_case( "slicing", all_set(v17[0:2] == vector( 1, 2 )) );
test_case( "slicing", all_set(v17[-3:-1] == vector( 3, 4 )) );
test_case( "slicing", all_set(v17[:2] == vector( 1, 2 )) );
test_case( "slicing", all_set(v17[2:] == vector( 3, 4, 5 )) );
}