Fix string enum bug (#79)

arr-ai · Mar 29, 2020 · 51d98e0 · 51d98e0
1 parent 009b9a1
commit 51d98e0
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 7 deletions.
diff --git a/rel/ops_set.go b/rel/ops_set.go
@@ -27,13 +27,13 @@ func NIntersect(a Set, bs ...Set) Set {
 func Union(a, b Set) Set {
 	if ga, ok := a.(GenericSet); ok {
 		if gb, ok := b.(GenericSet); ok {
-			return GenericSet{set: ga.set.Union(gb.set)}
+			return CanonicalSet(GenericSet{set: ga.set.Union(gb.set)})
 		}
 	}
 	for e := b.Enumerator(); e.MoveNext(); {
 		a = a.With(e.Current())
 	}
-	return a
+	return CanonicalSet(a)
 }
 
 func NUnion(sets ...Set) Set {

diff --git a/rel/value_set_generic.go b/rel/value_set_generic.go
@@ -70,6 +70,17 @@ func NewSet(values ...Value) Set {
 	return set
 }
 
+func CanonicalSet(s Set) Set {
+	if s, ok := s.(GenericSet); ok {
+		values := make([]Value, 0, s.Count())
+		for e := s.Enumerator(); e.MoveNext(); {
+			values = append(values, e.Current())
+		}
+		return NewSet(values...)
+	}
+	return s
+}
+
 // NewSetFrom constructs a genericSet from interfaces.
 func NewSetFrom(intfs ...interface{}) (Set, error) {
 	set := None

diff --git a/rel/value_set_str.go b/rel/value_set_str.go
@@ -200,14 +200,14 @@ func (s String) Map(f func(v Value) Value) Set {
 
 // Where returns a new String with all the Values satisfying predicate p.
 func (s String) Where(p func(v Value) bool) Set {
-	result := Set(s)
+	values := make([]Value, 0, s.Count())
 	for e := s.Enumerator(); e.MoveNext(); {
 		value := e.Current()
-		if !p(value) {
-			result = result.Without(value)
+		if p(value) {
+			values = append(values, value)
 		}
 	}
-	return result
+	return NewSet(values...)
 }
 
 // Call ...
@@ -250,7 +250,7 @@ func (e *stringValueEnumerator) MoveNext() bool {
 
 // Current returns the enumerator's current Value.
 func (e *stringValueEnumerator) Current() Value {
-	return NewStringCharTuple(e.i, e.s.s[e.i])
+	return NewStringCharTuple(e.s.offset+e.i, e.s.s[e.i])
 }
 
 type stringOffsetValueEnumerator struct {

diff --git a/syntax/expr_set_test.go b/syntax/expr_set_test.go
@@ -29,6 +29,10 @@ func TestArrayToString(t *testing.T) {
 	AssertCodeEvalsToType(t, rel.String{}, `[104, 101, 108, 108, 111] => (@:.@, @char:.@item)`)
 }
 
+func TestStringManipulation(t *testing.T) {
+	AssertCodesEvalToSameValue(t, `"Foo"`, `(\s //.str.upper(s where .@=0) | (s where .@>0))("foo")`)
+}
+
 func TestArray(t *testing.T) {
 	t.Parallel()
 	AssertCodesEvalToSameValue(t, `{|@,@item| (0, 1), (1, 2), (2, 3)}`, `[1, 2, 3]`)