Add custom iterator function support which enables implementing a REPL in jq

cmd/repl/main.go

@@ -0,0 +1,117 @@
+package main
+
+import (
+	"bufio"
+	"fmt"
+	"io"
+	"os"
+
+	"github.com/itchyny/gojq"
+)
+
+func read(c interface{}, a []interface{}) interface{} {
+	prompt, ok := a[0].(string)
+	if !ok {
+		return fmt.Errorf("%v: src is not a string", a[0])
+	}
+	fmt.Fprint(os.Stdout, prompt)
+	s := bufio.NewScanner(os.Stdin)
+	if ok = s.Scan(); !ok {
+		fmt.Fprintln(os.Stdout)
+		return io.EOF
+	}
+
+	return s.Text()
+}
+
+func eval(c interface{}, a []interface{}) interface{} {
+	src, ok := a[0].(string)
+	if !ok {
+		return fmt.Errorf("%v: src is not a string", a[0])
+	}
+	iter, err := replRun(c, src)
+	if err != nil {
+		return err
+	}
+
+	return iter
+}
+
+func print(c interface{}, a []interface{}) interface{} {
+	if _, err := fmt.Fprintln(os.Stdout, c); err != nil {
+		return err
+	}
+
+	return gojq.EmptyIter{}
+}
+
+func itertest(c interface{}, a []interface{}) interface{} {
+	return &gojq.SliceIter{Slice: []interface{}{1, 2, 3}}
+}
+
+func itererr(c interface{}, a []interface{}) interface{} {
+	return &gojq.SliceIter{Slice: []interface{}{1, fmt.Errorf("itervaluerr")}}
+}
+
+type preludeLoader struct{}
+
+func (preludeLoader) LoadInitModules() ([]*gojq.Query, error) {
+	replSrc := `
+def repl:
+	def _wrap: if (. | type) != "array" then [.] end;
+	def _repl:
+		try read("> ") as $e |
+		(try (.[] | eval($e)) catch . | print),
+		_repl;
+	_wrap | _repl;
+`
+	gq, err := gojq.Parse(replSrc)
+	if err != nil {
+		return nil, err
+	}
+
+	return []*gojq.Query{gq}, nil
+}
+
+func replRun(c interface{}, src string) (gojq.Iter, error) {
+	gq, err := gojq.Parse(src)
+	if err != nil {
+		return nil, err
+	}
+	gc, err := gojq.Compile(gq,
+		gojq.WithModuleLoader(preludeLoader{}),
+		gojq.WithFunction("read", 1, 1, read),
+		gojq.WithIterator("eval", 1, 1, eval),
+		gojq.WithIterator("print", 0, 0, print),
+
+		gojq.WithIterator("itertest", 0, 0, itertest),
+		gojq.WithIterator("itererr", 0, 0, itererr),
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	return gc.Run(c), nil
+}
+
+func main() {
+	expr := "repl"
+	if len(os.Args) > 1 {
+		expr = os.Args[1]
+	}
+	iter, err := replRun(nil, expr)
+	if err != nil {
+		panic(err)
+	}
+	for {
+		v, ok := iter.Next()
+		if !ok {
+			break
+		} else if err, ok := v.(error); ok {
+			fmt.Fprintf(os.Stderr, "err: %v\n", err)
+			break
+		} else if d, ok := v.([2]interface{}); ok {
+			fmt.Fprintf(os.Stdout, "%s: %v\n", d[0], d[1])
+		}
+	}
+}

compiler.go

@@ -897,6 +897,7 @@ func (c *compiler) compileFunc(e *Func) error {
 				e.Args,
 				nil,
 				false,
+				false,
 			)
 		case "input":
 			if c.inputIter == nil {
@@ -907,13 +908,15 @@ func (c *compiler) compileFunc(e *Func) error {
 				e.Args,
 				nil,
 				false,
+				false,
 			)
 		case "modulemeta":
 			return c.compileCallInternal(
 				[3]interface{}{c.funcModulemeta, 0, e.Name},
 				e.Args,
 				nil,
 				false,
+				false,
 			)
 		default:
 			return c.compileCall(e.Name, e.Args)
@@ -925,6 +928,7 @@ func (c *compiler) compileFunc(e *Func) error {
 			e.Args,
 			nil,
 			false,
+			fn.iterator,
 		)
 	}
 	return &funcNotFoundError{e}
@@ -1294,12 +1298,13 @@ func (c *compiler) compileCall(name string, args []*Query) error {
 		args,
 		nil,
 		name == "_index" || name == "_slice",
+		false,
 	)
 }
 
 func (c *compiler) compileCallPc(fn *funcinfo, args []*Query) error {
 	if len(args) == 0 {
-		return c.compileCallInternal(fn.pc, args, nil, false)
+		return c.compileCallInternal(fn.pc, args, nil, false, false)
 	}
 	xs, vars := make([]*Query, len(args)), make(map[int]bool, len(fn.args))
 	for i, j := range fn.argsorder {
@@ -1308,12 +1313,15 @@ func (c *compiler) compileCallPc(fn *funcinfo, args []*Query) error {
 			vars[i] = true
 		}
 	}
-	return c.compileCallInternal(fn.pc, xs, vars, false)
+	return c.compileCallInternal(fn.pc, xs, vars, false, false)
 }
 
-func (c *compiler) compileCallInternal(fn interface{}, args []*Query, vars map[int]bool, indexing bool) error {
+func (c *compiler) compileCallInternal(fn interface{}, args []*Query, vars map[int]bool, indexing bool, each bool) error {
 	if len(args) == 0 {
 		c.append(&code{op: opcall, v: fn})
+		if each {
+			c.append(&code{op: opeach})
+		}
 		return nil
 	}
 	idx := c.newVariable()
@@ -1367,6 +1375,9 @@ func (c *compiler) compileCallInternal(fn interface{}, args []*Query, vars map[i
 	}
 	c.append(&code{op: opload, v: idx})
 	c.append(&code{op: opcall, v: fn})
+	if each {
+		c.append(&code{op: opeach})
+	}
 	return nil
 }
 

execute.go

@@ -217,10 +217,17 @@ loop:
 				break loop
 			}
 			backtrack = false
+
+			var xc [2]interface{}
+			var xv interface{}
 			var xs [][2]interface{}
+
 			switch v := env.pop().(type) {
 			case [][2]interface{}:
-				xs = v
+				xc = v[0]
+				if len(v) > 1 {
+					xv = v[1:]
+				}
 			case []interface{}:
 				if !env.paths.empty() && (env.expdepth == 0 && !reflect.DeepEqual(v, env.paths.top().([2]interface{})[1])) {
 					err = &invalidPathIterError{v}
@@ -229,10 +236,14 @@ loop:
 				if len(v) == 0 {
 					break loop
 				}
-				xs = make([][2]interface{}, len(v))
+				xs := make([][2]interface{}, len(v))
 				for i, v := range v {
 					xs[i] = [2]interface{}{i, v}
 				}
+				xc = xs[0]
+				if len(xs) > 1 {
+					xv = xs[1:]
+				}
 			case map[string]interface{}:
 				if !env.paths.empty() && (env.expdepth == 0 && !reflect.DeepEqual(v, env.paths.top().([2]interface{})[1])) {
 					err = &invalidPathIterError{v}
@@ -250,21 +261,50 @@ loop:
 				sort.Slice(xs, func(i, j int) bool {
 					return xs[i][0].(string) < xs[j][0].(string)
 				})
+				xc = xs[0]
+				if len(xs) > 1 {
+					xv = xs[1:]
+				}
+			case Iter:
+				iv, ok := v.Next()
+				if !ok {
+					break loop
+				}
+				if e, ok := iv.(error); ok {
+					err = e
+					break loop
+				} else {
+					xc = [2]interface{}{0, iv}
+					xv = v
+				}
+
 			default:
 				err = &iteratorError{v}
 				break loop
 			}
-			if len(xs) > 1 {
-				env.push(xs[1:])
-				env.pushfork(code.op, pc)
-				env.pop()
-			}
-			env.push(xs[0][1])
-			if !env.paths.empty() {
-				if env.expdepth == 0 {
-					env.paths.push(xs[0])
+
+			// TODO: make debug nicer? backtrack correct?
+			if _, ok := xc[1].([2]interface{}); ok {
+				env.push(xv)
+				if !backtrack {
+					return xc[1], true
+				}
+				backtrack = false
+			} else {
+				if xv != nil {
+					env.push(xv)
+					env.pushfork(code.op, pc)
+					env.pop()
+				}
+
+				env.push(xc[1])
+				if !env.paths.empty() {
+					if env.expdepth == 0 {
+						env.paths.push(xc)
+					}
 				}
 			}
+
 		case opexpbegin:
 			env.expdepth++
 		case opexpend:

func.go

@@ -30,6 +30,7 @@ const (
 type function struct {
 	argcount int
 	callback func(interface{}, []interface{}) interface{}
+	iterator bool
 }
 
 func (fn function) accept(cnt int) bool {
@@ -60,7 +61,7 @@ func init() {
 		"contains":       argFunc1(funcContains),
 		"explode":        argFunc0(funcExplode),
 		"implode":        argFunc0(funcImplode),
-		"split":          {argcount1 | argcount2, funcSplit},
+		"split":          {argcount1 | argcount2, funcSplit, false},
 		"tojson":         argFunc0(funcToJSON),
 		"fromjson":       argFunc0(funcFromJSON),
 		"format":         argFunc1(funcFormat),
@@ -172,9 +173,9 @@ func init() {
 		"strptime":       argFunc1(funcStrptime),
 		"now":            argFunc0(funcNow),
 		"_match":         argFunc3(funcMatch),
-		"error":          {argcount0 | argcount1, funcError},
+		"error":          {argcount0 | argcount1, funcError, false},
 		"halt":           argFunc0(funcHalt),
-		"halt_error":     {argcount0 | argcount1, funcHaltError},
+		"halt_error":     {argcount0 | argcount1, funcHaltError, false},
 		"_type_error":    argFunc1(internalfuncTypeError),
 	}
 }
@@ -184,6 +185,7 @@ func argFunc0(fn func(interface{}) interface{}) function {
 		argcount0, func(v interface{}, _ []interface{}) interface{} {
 			return fn(v)
 		},
+		false,
 	}
 }
 
@@ -192,6 +194,7 @@ func argFunc1(fn func(interface{}, interface{}) interface{}) function {
 		argcount1, func(v interface{}, args []interface{}) interface{} {
 			return fn(v, args[0])
 		},
+		false,
 	}
 }
 
@@ -200,6 +203,7 @@ func argFunc2(fn func(interface{}, interface{}, interface{}) interface{}) functi
 		argcount2, func(v interface{}, args []interface{}) interface{} {
 			return fn(v, args[0], args[1])
 		},
+		false,
 	}
 }
 
@@ -208,6 +212,7 @@ func argFunc3(fn func(interface{}, interface{}, interface{}, interface{}) interf
 		argcount3, func(v interface{}, args []interface{}) interface{} {
 			return fn(v, args[0], args[1], args[2])
 		},
+		false,
 	}
 }
 

iter.go

@@ -21,3 +21,28 @@ func (c *unitIter) Next() (interface{}, bool) {
 	}
 	return nil, false
 }
+
+// SliceIter is a Iter that iterate a interface{} slice from first to last value
+type SliceIter struct{ Slice []interface{} }
+
+// Next value in slice or no value if at end
+func (i *SliceIter) Next() (interface{}, bool) {
+	if len(i.Slice) == 0 {
+		return nil, false
+	}
+	e := i.Slice[0]
+	i.Slice = i.Slice[1:]
+	return e, true
+}
+
+// EmptyIter is a Iter that return no value, similar to "empty" keyword.
+type EmptyIter struct{}
+
+// Next returns no value
+func (EmptyIter) Next() (interface{}, bool) { return nil, false }
+
+// IterFn is a Iter that calls a provided next function
+type IterFn func() (interface{}, bool)
+
+// Next value in slice or no value if at end
+func (i IterFn) Next() (interface{}, bool) { return i() }