query.go

package gojq

import (
	"cmp"
	"context"
	"strings"
)

// Parse a query string, and returns the query struct.
//
// If parsing failed, it returns an error of type [*ParseError], which has
// the byte offset and the invalid token. The byte offset is the scanned bytes
// when the error occurred. The token is empty if the error occurred after
// scanning the entire query string.
func Parse(src string) (*Query, error) {
	l := newLexer(src)
	if yyParse(l) > 0 {
		return nil, l.err
	}
	return l.result, nil
}

// Query represents the abstract syntax tree of a jq query.
type Query struct {
	Meta     *ConstObject
	Imports  []*Import
	FuncDefs []*FuncDef
	Term     *Term
	Left     *Query
	Right    *Query
	Patterns []*Pattern
	Op       Operator
}

// Run the query.
//
// It is safe to call this method in goroutines, to reuse a parsed [*Query].
func (e *Query) Run(v any) Iter {
	return e.RunWithContext(context.Background(), v)
}

// RunWithContext runs the query with context.
func (e *Query) RunWithContext(ctx context.Context, v any) Iter {
	code, err := Compile(e)
	if err != nil {
		return NewIter(err)
	}
	return code.RunWithContext(ctx, v)
}

func (e *Query) String() string {
	var s strings.Builder
	e.writeTo(&s)
	return s.String()
}

func (e *Query) writeTo(s *strings.Builder) {
	if e.Meta != nil {
		s.WriteString("module ")
		e.Meta.writeTo(s)
		s.WriteString(";\n")
	}
	for _, im := range e.Imports {
		im.writeTo(s)
	}
	for _, fd := range e.FuncDefs {
		fd.writeTo(s)
		s.WriteByte(' ')
	}
	if e.Term != nil {
		e.Term.writeTo(s)
	} else if e.Right != nil {
		e.Left.writeTo(s)
		writePatternsTo(s, e.Patterns)
		if e.Op != OpComma {
			s.WriteByte(' ')
		}
		s.WriteString(e.Op.String())
		s.WriteByte(' ')
		e.Right.writeTo(s)
	}
}

func (e *Query) toIndexKey() any {
	if e.Term == nil {
		return nil
	}
	return e.Term.toIndexKey()
}

func (e *Query) toIndices(xs []any) []any {
	if e.Term == nil {
		return nil
	}
	return e.Term.toIndices(xs)
}

// Import ...
type Import struct {
	ImportPath  string
	ImportAlias string
	IncludePath string
	Meta        *ConstObject
}

func (e *Import) String() string {
	var s strings.Builder
	e.writeTo(&s)
	return s.String()
}

func (e *Import) writeTo(s *strings.Builder) {
	if e.ImportPath != "" {
		s.WriteString("import ")
		jsonEncodeString(s, e.ImportPath)
		s.WriteString(" as ")
		s.WriteString(e.ImportAlias)
	} else {
		s.WriteString("include ")
		jsonEncodeString(s, e.IncludePath)
	}
	if e.Meta != nil {
		s.WriteByte(' ')
		e.Meta.writeTo(s)
	}
	s.WriteString(";\n")
}

// FuncDef ...
type FuncDef struct {
	Name string
	Args []string
	Body *Query
}

func (e *FuncDef) String() string {
	var s strings.Builder
	e.writeTo(&s)
	return s.String()
}

func (e *FuncDef) writeTo(s *strings.Builder) {
	s.WriteString("def ")
	s.WriteString(e.Name)
	if len(e.Args) > 0 {
		s.WriteByte('(')
		for i, e := range e.Args {
			if i > 0 {
				s.WriteString("; ")
			}
			s.WriteString(e)
		}
		s.WriteByte(')')
	}
	s.WriteString(": ")
	e.Body.writeTo(s)
	s.WriteByte(';')
}

// Term ...
type Term struct {
	Type       TermType
	Index      *Index
	Func       *Func
	Object     *Object
	Array      *Array
	Number     string
	Unary      *Unary
	Format     string
	Str        *String
	If         *If
	Try        *Try
	Reduce     *Reduce
	Foreach    *Foreach
	Label      *Label
	Break      string
	Query      *Query
	SuffixList []*Suffix
}

func (e *Term) String() string {
	var s strings.Builder
	e.writeTo(&s)
	return s.String()
}

func (e *Term) writeTo(s *strings.Builder) {
	switch e.Type {
	case TermTypeIdentity:
		s.WriteByte('.')
	case TermTypeRecurse:
		s.WriteString("..")
	case TermTypeNull:
		s.WriteString("null")
	case TermTypeTrue:
		s.WriteString("true")
	case TermTypeFalse:
		s.WriteString("false")
	case TermTypeIndex:
		e.Index.writeTo(s)
	case TermTypeFunc:
		e.Func.writeTo(s)
	case TermTypeObject:
		e.Object.writeTo(s)
	case TermTypeArray:
		e.Array.writeTo(s)
	case TermTypeNumber:
		s.WriteString(e.Number)
	case TermTypeUnary:
		e.Unary.writeTo(s)
	case TermTypeFormat:
		s.WriteString(e.Format)
		if e.Str != nil {
			s.WriteByte(' ')
			e.Str.writeTo(s)
		}
	case TermTypeString:
		e.Str.writeTo(s)
	case TermTypeIf:
		e.If.writeTo(s)
	case TermTypeTry:
		e.Try.writeTo(s)
	case TermTypeReduce:
		e.Reduce.writeTo(s)
	case TermTypeForeach:
		e.Foreach.writeTo(s)
	case TermTypeLabel:
		e.Label.writeTo(s)
	case TermTypeBreak:
		s.WriteString("break ")
		s.WriteString(e.Break)
	case TermTypeQuery:
		s.WriteByte('(')
		e.Query.writeTo(s)
		s.WriteByte(')')
	}
	for _, e := range e.SuffixList {
		e.writeTo(s)
	}
}

func (e *Term) toIndexKey() any {
	if len(e.SuffixList) > 0 {
		return nil
	}
	switch e.Type {
	case TermTypeNumber:
		return toNumber(e.Number)
	case TermTypeUnary:
		return e.Unary.toNumber()
	case TermTypeString:
		if e.Str.Queries == nil {
			return e.Str.Str
		}
		return nil
	default:
		return nil
	}
}

func (e *Term) toIndices(xs []any) []any {
	switch e.Type {
	case TermTypeIndex:
		if xs = e.Index.toIndices(xs); xs == nil {
			return nil
		}
	case TermTypeQuery:
		if xs = e.Query.toIndices(xs); xs == nil {
			return nil
		}
	default:
		return nil
	}
	for _, s := range e.SuffixList {
		if xs = s.toIndices(xs); xs == nil {
			return nil
		}
	}
	return xs
}

func (e *Term) toNumber() any {
	if e.Type == TermTypeNumber && len(e.SuffixList) == 0 {
		return toNumber(e.Number)
	}
	return nil
}

// Unary ...
type Unary struct {
	Op   Operator
	Term *Term
}

func (e *Unary) String() string {
	var s strings.Builder
	e.writeTo(&s)
	return s.String()
}

func (e *Unary) writeTo(s *strings.Builder) {
	s.WriteString(e.Op.String())
	e.Term.writeTo(s)
}

func (e *Unary) toNumber() any {
	v := e.Term.toNumber()
	if v != nil && e.Op == OpSub {
		v = funcOpNegate(v)
	}
	return v
}

// Pattern ...
type Pattern struct {
	Name   string
	Array  []*Pattern
	Object []*PatternObject
}

func (e *Pattern) String() string {
	var s strings.Builder
	e.writeTo(&s)
	return s.String()
}

func (e *Pattern) writeTo(s *strings.Builder) {
	if e.Name != "" {
		s.WriteString(e.Name)
	} else if len(e.Array) > 0 {
		s.WriteByte('[')
		for i, e := range e.Array {
			if i > 0 {
				s.WriteString(", ")
			}
			e.writeTo(s)
		}
		s.WriteByte(']')
	} else if len(e.Object) > 0 {
		s.WriteByte('{')
		for i, e := range e.Object {
			if i > 0 {
				s.WriteString(", ")
			}
			e.writeTo(s)
		}
		s.WriteByte('}')
	}
}

func writePatternsTo(s *strings.Builder, ps []*Pattern) {
	for i, p := range ps {
		if i == 0 {
			s.WriteString(" as ")
		} else {
			s.WriteString(" ?// ")
		}
		p.writeTo(s)
	}
}

// PatternObject ...
type PatternObject struct {
	Key       string
	KeyString *String
	KeyQuery  *Query
	Val       *Pattern
}

func (e *PatternObject) String() string {
	var s strings.Builder
	e.writeTo(&s)
	return s.String()
}

func (e *PatternObject) writeTo(s *strings.Builder) {
	if e.Key != "" {
		s.WriteString(e.Key)
	} else if e.KeyString != nil {
		e.KeyString.writeTo(s)
	} else if e.KeyQuery != nil {
		s.WriteByte('(')
		e.KeyQuery.writeTo(s)
		s.WriteByte(')')
	}
	if e.Val != nil {
		s.WriteString(": ")
		e.Val.writeTo(s)
	}
}

// Index ...
type Index struct {
	Name    string
	Str     *String
	Start   *Query
	End     *Query
	IsSlice bool
}

func (e *Index) String() string {
	var s strings.Builder
	e.writeTo(&s)
	return s.String()
}

func (e *Index) writeTo(s *strings.Builder) {
	if l := s.Len(); l > 0 {
		// ". .x" != "..x" and "0 .x" != "0.x"
		if c := s.String()[l-1]; c == '.' || '0' <= c && c <= '9' {
			s.WriteByte(' ')
		}
	}
	s.WriteByte('.')
	e.writeSuffixTo(s)
}

func (e *Index) writeSuffixTo(s *strings.Builder) {
	if e.Name != "" {
		s.WriteString(e.Name)
	} else if e.Str != nil {
		e.Str.writeTo(s)
	} else {
		s.WriteByte('[')
		if e.IsSlice {
			if e.Start != nil {
				e.Start.writeTo(s)
			}
			s.WriteByte(':')
			if e.End != nil {
				e.End.writeTo(s)
			}
		} else {
			e.Start.writeTo(s)
		}
		s.WriteByte(']')
	}
}

func (e *Index) toIndexKey() any {
	if e.Name != "" {
		return e.Name
	} else if e.Str != nil {
		if e.Str.Queries == nil {
			return e.Str.Str
		}
	} else if !e.IsSlice {
		return e.Start.toIndexKey()
	} else {
		var start, end any
		ok := true
		if e.Start != nil {
			start = e.Start.toIndexKey()
			ok = start != nil
		}
		if e.End != nil && ok {
			end = e.End.toIndexKey()
			ok = end != nil
		}
		if ok {
			return map[string]any{"start": start, "end": end}
		}
	}
	return nil
}

func (e *Index) toIndices(xs []any) []any {
	if k := e.toIndexKey(); k != nil {
		return append(xs, k)
	}
	return nil
}

// Func ...
type Func struct {
	Name string
	Args []*Query
}

func (e *Func) String() string {
	var s strings.Builder
	e.writeTo(&s)
	return s.String()
}

func (e *Func) writeTo(s *strings.Builder) {
	s.WriteString(e.Name)
	if len(e.Args) > 0 {
		s.WriteByte('(')
		for i, e := range e.Args {
			if i > 0 {
				s.WriteString("; ")
			}
			e.writeTo(s)
		}
		s.WriteByte(')')
	}
}

// String ...
type String struct {
	Str     string
	Queries []*Query
}

func (e *String) String() string {
	var s strings.Builder
	e.writeTo(&s)
	return s.String()
}

func (e *String) writeTo(s *strings.Builder) {
	if e.Queries == nil {
		jsonEncodeString(s, e.Str)
		return
	}
	s.WriteByte('"')
	for _, e := range e.Queries {
		if e.Term.Str == nil {
			s.WriteString(`\`)
			e.writeTo(s)
		} else {
			es := e.String()
			s.WriteString(es[1 : len(es)-1])
		}
	}
	s.WriteByte('"')
}

// Object ...
type Object struct {
	KeyVals []*ObjectKeyVal
}

func (e *Object) String() string {
	var s strings.Builder
	e.writeTo(&s)
	return s.String()
}

func (e *Object) writeTo(s *strings.Builder) {
	if len(e.KeyVals) == 0 {
		s.WriteString("{}")
		return
	}
	s.WriteString("{ ")
	for i, kv := range e.KeyVals {
		if i > 0 {
			s.WriteString(", ")
		}
		kv.writeTo(s)
	}
	s.WriteString(" }")
}

// ObjectKeyVal ...
type ObjectKeyVal struct {
	Key       string
	KeyString *String
	KeyQuery  *Query
	Val       *Query
}

func (e *ObjectKeyVal) String() string {
	var s strings.Builder
	e.writeTo(&s)
	return s.String()
}

func (e *ObjectKeyVal) writeTo(s *strings.Builder) {
	if e.Key != "" {
		s.WriteString(e.Key)
	} else if e.KeyString != nil {
		e.KeyString.writeTo(s)
	} else if e.KeyQuery != nil {
		s.WriteByte('(')
		e.KeyQuery.writeTo(s)
		s.WriteByte(')')
	}
	if e.Val != nil {
		s.WriteString(": ")
		e.Val.writeTo(s)
	}
}

// Array ...
type Array struct {
	Query *Query
}

func (e *Array) String() string {
	var s strings.Builder
	e.writeTo(&s)
	return s.String()
}

func (e *Array) writeTo(s *strings.Builder) {
	s.WriteByte('[')
	if e.Query != nil {
		e.Query.writeTo(s)
	}
	s.WriteByte(']')
}

// Suffix ...
type Suffix struct {
	Index    *Index
	Iter     bool
	Optional bool
}

func (e *Suffix) String() string {
	var s strings.Builder
	e.writeTo(&s)
	return s.String()
}

func (e *Suffix) writeTo(s *strings.Builder) {
	if e.Index != nil {
		if e.Index.Name != "" || e.Index.Str != nil {
			e.Index.writeTo(s)
		} else {
			e.Index.writeSuffixTo(s)
		}
	} else if e.Iter {
		s.WriteString("[]")
	} else if e.Optional {
		s.WriteByte('?')
	}
}

func (e *Suffix) toTerm() *Term {
	if e.Index != nil {
		return &Term{Type: TermTypeIndex, Index: e.Index}
	} else if e.Iter {
		return &Term{Type: TermTypeIdentity, SuffixList: []*Suffix{{Iter: true}}}
	} else {
		return nil
	}
}

func (e *Suffix) toIndices(xs []any) []any {
	if e.Index == nil {
		return nil
	}
	return e.Index.toIndices(xs)
}

// If ...
type If struct {
	Cond *Query
	Then *Query
	Elif []*IfElif
	Else *Query
}

func (e *If) String() string {
	var s strings.Builder
	e.writeTo(&s)
	return s.String()
}

func (e *If) writeTo(s *strings.Builder) {
	s.WriteString("if ")
	e.Cond.writeTo(s)
	s.WriteString(" then ")
	e.Then.writeTo(s)
	for _, e := range e.Elif {
		s.WriteByte(' ')
		e.writeTo(s)
	}
	if e.Else != nil {
		s.WriteString(" else ")
		e.Else.writeTo(s)
	}
	s.WriteString(" end")
}

// IfElif ...
type IfElif struct {
	Cond *Query
	Then *Query
}

func (e *IfElif) String() string {
	var s strings.Builder
	e.writeTo(&s)
	return s.String()
}

func (e *IfElif) writeTo(s *strings.Builder) {
	s.WriteString("elif ")
	e.Cond.writeTo(s)
	s.WriteString(" then ")
	e.Then.writeTo(s)
}

// Try ...
type Try struct {
	Body  *Query
	Catch *Query
}

func (e *Try) String() string {
	var s strings.Builder
	e.writeTo(&s)
	return s.String()
}

func (e *Try) writeTo(s *strings.Builder) {
	s.WriteString("try ")
	e.Body.writeTo(s)
	if e.Catch != nil {
		s.WriteString(" catch ")
		e.Catch.writeTo(s)
	}
}

// Reduce ...
type Reduce struct {
	Query    *Query
	Patterns []*Pattern
	Start    *Query
	Update   *Query
}

func (e *Reduce) String() string {
	var s strings.Builder
	e.writeTo(&s)
	return s.String()
}

func (e *Reduce) writeTo(s *strings.Builder) {
	s.WriteString("reduce ")
	e.Query.writeTo(s)
	writePatternsTo(s, e.Patterns)
	s.WriteString(" (")
	e.Start.writeTo(s)
	s.WriteString("; ")
	e.Update.writeTo(s)
	s.WriteByte(')')
}

// Foreach ...
type Foreach struct {
	Query    *Query
	Patterns []*Pattern
	Start    *Query
	Update   *Query
	Extract  *Query
}

func (e *Foreach) String() string {
	var s strings.Builder
	e.writeTo(&s)
	return s.String()
}

func (e *Foreach) writeTo(s *strings.Builder) {
	s.WriteString("foreach ")
	e.Query.writeTo(s)
	writePatternsTo(s, e.Patterns)
	s.WriteString(" (")
	e.Start.writeTo(s)
	s.WriteString("; ")
	e.Update.writeTo(s)
	if e.Extract != nil {
		s.WriteString("; ")
		e.Extract.writeTo(s)
	}
	s.WriteByte(')')
}

// Label ...
type Label struct {
	Ident string
	Body  *Query
}

func (e *Label) String() string {
	var s strings.Builder
	e.writeTo(&s)
	return s.String()
}

func (e *Label) writeTo(s *strings.Builder) {
	s.WriteString("label ")
	s.WriteString(e.Ident)
	s.WriteString(" | ")
	e.Body.writeTo(s)
}

// ConstTerm ...
type ConstTerm struct {
	Object *ConstObject
	Array  *ConstArray
	Number string
	Str    string
	Null   bool
	True   bool
	False  bool
}

func (e *ConstTerm) String() string {
	var s strings.Builder
	e.writeTo(&s)
	return s.String()
}

func (e *ConstTerm) writeTo(s *strings.Builder) {
	if e.Object != nil {
		e.Object.writeTo(s)
	} else if e.Array != nil {
		e.Array.writeTo(s)
	} else if e.Number != "" {
		s.WriteString(e.Number)
	} else if e.Null {
		s.WriteString("null")
	} else if e.True {
		s.WriteString("true")
	} else if e.False {
		s.WriteString("false")
	} else {
		jsonEncodeString(s, e.Str)
	}
}

func (e *ConstTerm) toValue() any {
	if e.Object != nil {
		return e.Object.ToValue()
	} else if e.Array != nil {
		return e.Array.toValue()
	} else if e.Number != "" {
		return toNumber(e.Number)
	} else if e.Null {
		return nil
	} else if e.True {
		return true
	} else if e.False {
		return false
	} else {
		return e.Str
	}
}

func (e *ConstTerm) toString() (string, bool) {
	if e.Object != nil || e.Array != nil ||
		e.Number != "" || e.Null || e.True || e.False {
		return "", false
	}
	return e.Str, true
}

// ConstObject ...
type ConstObject struct {
	KeyVals []*ConstObjectKeyVal
}

func (e *ConstObject) String() string {
	var s strings.Builder
	e.writeTo(&s)
	return s.String()
}

func (e *ConstObject) writeTo(s *strings.Builder) {
	if len(e.KeyVals) == 0 {
		s.WriteString("{}")
		return
	}
	s.WriteString("{ ")
	for i, kv := range e.KeyVals {
		if i > 0 {
			s.WriteString(", ")
		}
		kv.writeTo(s)
	}
	s.WriteString(" }")
}

// ToValue converts the object to map[string]any.
func (e *ConstObject) ToValue() map[string]any {
	if e == nil {
		return nil
	}
	v := make(map[string]any, len(e.KeyVals))
	for _, e := range e.KeyVals {
		v[cmp.Or(e.Key, e.KeyString)] = e.Val.toValue()
	}
	return v
}

// ConstObjectKeyVal ...
type ConstObjectKeyVal struct {
	Key       string
	KeyString string
	Val       *ConstTerm
}

func (e *ConstObjectKeyVal) String() string {
	var s strings.Builder
	e.writeTo(&s)
	return s.String()
}

func (e *ConstObjectKeyVal) writeTo(s *strings.Builder) {
	if e.Key != "" {
		s.WriteString(e.Key)
	} else {
		jsonEncodeString(s, e.KeyString)
	}
	s.WriteString(": ")
	e.Val.writeTo(s)
}

// ConstArray ...
type ConstArray struct {
	Elems []*ConstTerm
}

func (e *ConstArray) String() string {
	var s strings.Builder
	e.writeTo(&s)
	return s.String()
}

func (e *ConstArray) writeTo(s *strings.Builder) {
	s.WriteByte('[')
	for i, e := range e.Elems {
		if i > 0 {
			s.WriteString(", ")
		}
		e.writeTo(s)
	}
	s.WriteByte(']')
}

func (e *ConstArray) toValue() []any {
	v := make([]any, len(e.Elems))
	for i, e := range e.Elems {
		v[i] = e.toValue()
	}
	return v
}