put: Assignment to array or set

Update the put operator to allow assignments to elements in an array or
set.

Closes #4798

Author: mattnibs

compiler/semantic/op.go

@@ -619,16 +619,6 @@ func (a *analyzer) semOp(o ast.Op, seq dag.Seq) (dag.Seq, error) {
 		if err != nil {
 			return nil, err
 		}
-		// We can do collision checking on static paths, so check what we can.
-		var fields field.List
-		for _, a := range assignments {
-			if this, ok := a.LHS.(*dag.This); ok {
-				fields = append(fields, this.Path)
-			}
-		}
-		if err := expr.CheckPutFields(fields); err != nil {
-			return nil, fmt.Errorf("put: %w", err)
-		}
 		return append(seq, &dag.Put{
 			Kind: "Put",
 			Args: assignments,

docs/language/operators/put.md

@@ -82,5 +82,5 @@ echo '{a:1} 1' | zq -z 'b:=2' -
 =>
 ```mdtest-output
 {a:1,b:2}
-error({message:"put: not a record",on:1})
+error({message:"put: not a puttable element",on:1})
 ```

runtime/expr/cutter.go

@@ -94,11 +94,11 @@ func (c *Cutter) Eval(ectx Context, in *zed.Value) *zed.Value {
 func (c *Cutter) lookupBuilder(ectx Context, in *zed.Value) (*recordBuilderCachedTypes, field.List, error) {
 	paths := c.fieldRefs[:0]
 	for _, p := range c.lvals {
-		path, err := p.Eval(ectx, in)
+		path, err := p.EvalAsRecordPath(ectx, in)
 		if err != nil {
 			return nil, nil, err
 		}
-		if path.IsEmpty() {
+		if len(path) == 0 {
 			return nil, nil, errors.New("'this' not allowed (use record literal)")
 		}
 		paths = append(paths, path)

runtime/expr/dynfield/path.go

@@ -0,0 +1,41 @@
+package dynfield
+
+import (
+	"github.com/brimdata/zed"
+	"github.com/brimdata/zed/zson"
+)
+
+type Path []zed.Value
+
+func (p Path) Append(b []byte) []byte {
+	for i, v := range p {
+		if i > 0 {
+			b = append(b, 0)
+		}
+		b = append(b, v.Bytes()...)
+	}
+	return b
+}
+
+func (p Path) String() string {
+	var b []byte
+	for i, v := range p {
+		if i > 0 {
+			b = append(b, '.')
+		}
+		b = append(b, zson.FormatValue(&v)...)
+	}
+	return string(b)
+}
+
+type List []Path
+
+func (l List) Append(b []byte) []byte {
+	for i, path := range l {
+		if i > 0 {
+			b = append(b, ',')
+		}
+		b = path.Append(b)
+	}
+	return b
+}

runtime/expr/lval.go

@@ -5,12 +5,14 @@ import (
 
 	"github.com/brimdata/zed"
 	"github.com/brimdata/zed/pkg/field"
+	"github.com/brimdata/zed/runtime/expr/dynfield"
 	"github.com/brimdata/zed/zson"
 )
 
 type Lval struct {
-	Elems []LvalElem
-	cache field.Path
+	Elems      []LvalElem
+	cache      []zed.Value
+	fieldCache field.Path
 }
 
 func NewLval(evals []LvalElem) *Lval {
@@ -19,18 +21,36 @@ func NewLval(evals []LvalElem) *Lval {
 
 // Eval returns the path of the lval. If there's an error the returned *zed.Value
 // will not be nill.
-func (l *Lval) Eval(ectx Context, this *zed.Value) (field.Path, error) {
+func (l *Lval) Eval(ectx Context, this *zed.Value) (dynfield.Path, error) {
 	l.cache = l.cache[:0]
 	for _, e := range l.Elems {
-		name, err := e.Eval(ectx, this)
+		val, err := e.Eval(ectx, this)
 		if err != nil {
 			return nil, err
 		}
-		l.cache = append(l.cache, name)
+		l.cache = append(l.cache, *val)
 	}
 	return l.cache, nil
 }
 
+func (l *Lval) EvalAsRecordPath(ectx Context, this *zed.Value) (field.Path, error) {
+	l.fieldCache = l.fieldCache[:0]
+	for _, e := range l.Elems {
+		val, err := e.Eval(ectx, this)
+		if err != nil {
+			return nil, err
+		}
+		if !val.IsString() {
+			// XXX Add context to error so we know what element is failing but
+			// let's wait until we can test this so we have a feel for what we
+			// want to see.
+			return nil, errors.New("field reference is not a string")
+		}
+		l.fieldCache = append(l.fieldCache, val.AsString())
+	}
+	return l.fieldCache, nil
+}
+
 // Path returns the receiver's path.  Path returns false when the receiver
 // contains a dynamic element.
 func (l *Lval) Path() (field.Path, bool) {
@@ -46,15 +66,15 @@ func (l *Lval) Path() (field.Path, bool) {
 }
 
 type LvalElem interface {
-	Eval(ectx Context, this *zed.Value) (string, error)
+	Eval(ectx Context, this *zed.Value) (*zed.Value, error)
 }
 
 type StaticLvalElem struct {
 	Name string
 }
 
-func (l *StaticLvalElem) Eval(_ Context, _ *zed.Value) (string, error) {
-	return l.Name, nil
+func (l *StaticLvalElem) Eval(_ Context, _ *zed.Value) (*zed.Value, error) {
+	return zed.NewString(l.Name), nil
 }
 
 type ExprLvalElem struct {
@@ -69,17 +89,12 @@ func NewExprLvalElem(zctx *zed.Context, e Evaluator) *ExprLvalElem {
 	}
 }
 
-func (l *ExprLvalElem) Eval(ectx Context, this *zed.Value) (string, error) {
+func (l *ExprLvalElem) Eval(ectx Context, this *zed.Value) (*zed.Value, error) {
 	val := l.eval.Eval(ectx, this)
 	if val.IsError() {
-		return "", lvalErr(ectx, val)
-	}
-	if !val.IsString() {
-		if val = l.caster.Eval(ectx, val); val.IsError() {
-			return "", errors.New("field reference is not a string")
-		}
+		return nil, lvalErr(ectx, val)
 	}
-	return val.AsString(), nil
+	return val, nil
 }
 
 func lvalErr(ectx Context, errVal *zed.Value) error {

runtime/expr/pathbuilder/builder.go

@@ -0,0 +1,97 @@
+package pathbuilder
+
+import (
+	"errors"
+
+	"github.com/brimdata/zed"
+	"github.com/brimdata/zed/runtime/expr/dynfield"
+)
+
+type builder struct {
+	inputCount int
+	base       Step
+}
+
+func New(base zed.Type, paths []dynfield.Path, leafs []zed.Value) (Step, error) {
+	if len(paths) != len(leafs) {
+		return nil, errors.New("paths and leafs must be the same length")
+	}
+	b := &builder{base: newLeafStep(base, -1)}
+	for i, p := range paths {
+		if err := b.Put(p, leafs[i].Type); err != nil {
+			return nil, err
+		}
+	}
+	return b.base, nil
+}
+
+func (m *builder) Put(p dynfield.Path, leaf zed.Type) error {
+	defer func() { m.inputCount++ }()
+	return m.put(&m.base, p, leaf)
+}
+
+func (m *builder) put(parent *Step, p dynfield.Path, typ zed.Type) error {
+	// Actually let's do this differently. If current is a string then we are
+	// putting to a record. When we support maps we'll need to check for that.
+	if p[0].IsString() {
+		return m.putRecord(parent, p, typ)
+	}
+	// This could be for a map or a set but keep it simple for now.
+	if zed.IsInteger(p[0].Type.ID()) {
+		return m.putVector(parent, p, typ)
+	}
+	// if zed.TypeUnder(parent.typeof())
+	return errors.New("unsupported types")
+}
+
+func (m *builder) putRecord(s *Step, p dynfield.Path, typ zed.Type) error {
+	current, p := p[0], p[1:]
+	rstep, ok := (*s).(*recordStep)
+	if !ok {
+		// If this is a leafStep with a type of record than we need to
+		// initialize a recordStep with fields, otherwise just replace this will
+		// a recordStep.
+		var fields []zed.Field
+		if lstep, ok := (*s).(*leafStep); ok && zed.TypeRecordOf(lstep.typ) != nil {
+			fields = zed.TypeRecordOf(lstep.typ).Fields
+		}
+		rstep = newRecordStep(fields)
+		if *s == m.base {
+			rstep.isBase = true
+		}
+		*s = rstep
+	}
+	i := rstep.lookup(current.AsString())
+	field := &rstep.fields[i]
+	if len(p) == 0 {
+		field.step = newLeafStep(typ, m.inputCount)
+		return nil
+	}
+	return m.put(&field.step, p, typ)
+}
+
+func (m *builder) putVector(s *Step, p dynfield.Path, typ zed.Type) error {
+	current, p := p[0], p[1:]
+	vstep, ok := (*s).(*vectorStep)
+	if !ok {
+		// If this is a leafStep with a type of array than we need to
+		// initialize a arrayStep with fields, otherwise just replace this with
+		// an arrayStep.
+		vstep = &vectorStep{}
+		if lstep, ok := (*s).(*leafStep); ok && zed.InnerType(lstep.typ) != nil {
+			vstep.inner = zed.InnerType(lstep.typ)
+			_, vstep.isSet = zed.TypeUnder(lstep.typ).(*zed.TypeSet)
+		}
+		if *s == m.base {
+			vstep.isBase = true
+		}
+		*s = vstep
+	}
+	at := vstep.lookup(int(current.AsInt()))
+	elem := &vstep.elems[at]
+	if len(p) == 0 {
+		elem.step = newLeafStep(typ, m.inputCount)
+		return nil
+	}
+	return m.put(&elem.step, p, typ)
+}

runtime/expr/pathbuilder/builder_test.go

@@ -0,0 +1,108 @@
+package pathbuilder
+
+import (
+	"testing"
+
+	"github.com/brimdata/zed"
+	"github.com/brimdata/zed/runtime/expr/dynfield"
+	"github.com/brimdata/zed/zcode"
+	"github.com/brimdata/zed/zson"
+	"github.com/stretchr/testify/require"
+)
+
+func parsePath(zctx *zed.Context, ss ...string) dynfield.Path {
+	var path dynfield.Path
+	for _, s := range ss {
+		path = append(path, *zson.MustParseValue(zctx, s))
+	}
+	return path
+}
+
+type testCase struct {
+	describe string
+	base     string
+	paths    [][]string
+	values   []string
+	expected string
+}
+
+func runTestCase(t *testing.T, c testCase) {
+	zctx := zed.NewContext()
+	var baseTyp zed.Type
+	var baseBytes []byte
+	if c.base != "" {
+		base := zson.MustParseValue(zctx, c.base)
+		baseTyp, baseBytes = base.Type, base.Bytes()
+	}
+	var paths []dynfield.Path
+	for _, ss := range c.paths {
+		paths = append(paths, parsePath(zctx, ss...))
+	}
+	var values []zed.Value
+	for _, s := range c.values {
+		values = append(values, *zson.MustParseValue(zctx, s))
+	}
+	step, err := New(baseTyp, paths, values)
+	require.NoError(t, err)
+	var b zcode.Builder
+	typ, err := step.Build(zctx, &b, baseBytes, values)
+	require.NoError(t, err)
+	val := zed.NewValue(typ, b.Bytes())
+	require.Equal(t, c.expected, zson.FormatValue(val))
+}
+
+func TestIt(t *testing.T) {
+	runTestCase(t, testCase{
+		base: `{"a": 1, "b": 2}`,
+		paths: [][]string{
+			{`"c"`, `"a"`, `"a"`},
+			{`"c"`, `"b"`},
+			{`"c"`, `"c"`},
+		},
+		values: []string{
+			`45`,
+			`"string"`,
+			"127.0.0.1",
+		},
+		expected: `{a:1,b:2,c:{a:{a:45},b:"string",c:127.0.0.1}}`,
+	})
+	runTestCase(t, testCase{
+		base: `{"a": [1,{foo:"bar"}]}`,
+		paths: [][]string{
+			{`"a"`, `0`},
+			{`"a"`, `1`, `"foo"`},
+		},
+		values: []string{
+			`"hi"`,
+			`"baz"`,
+		},
+		expected: `{a:["hi",{foo:"baz"}]}`,
+	})
+	runTestCase(t, testCase{
+		describe: "create from empty base",
+		paths: [][]string{
+			{`"a"`},
+			{`"b"`},
+		},
+		values: []string{
+			`"foo"`,
+			`"bar"`,
+		},
+		expected: `{a:"foo",b:"bar"}`,
+	})
+	runTestCase(t, testCase{
+		describe: "assign to base level array",
+		base:     `["a", "b", "c"]`,
+		paths: [][]string{
+			{`0`},
+			{`1`},
+			{`2`},
+		},
+		values: []string{
+			`"foo"`,
+			`"bar"`,
+			`"baz"`,
+		},
+		expected: `["foo","bar","baz"]`,
+	})
+}