improve tail call optimization (close #86)

Previously it optimizes only functions with no local variables, like _while and
_until, using opjump instruction. This commit also optimizes functions with
local variables but with no arguments and there is no following fork point
because local variables may be referenced after backtracking. For tail call
with following fork points, it now uses new opcallrec instruction to
optimize return address but allocates for local variables.

code.go

@@ -26,6 +26,7 @@ const (
  opjump
  opjumpifnot
  opcall
+ opcallrec
  oppushpc
  opcallpc
  opscope
@@ -75,6 +76,8 @@ func (op opcode) String() string {
  return "jumpifnot"
  case opcall:
  return "call"
+ case opcallrec:
+ return "callrec"
  case oppushpc:
  return "pushpc"
  case opcallpc:

compiler.go

@@ -96,7 +96,7 @@ func Compile(q *Query, options ...CompilerOption) (*Code, error) {
  scope := c.newScope()
  c.scopes = []*scopeinfo{scope}
  defer c.lazy(func() *code {
- return &code{op: opscope, v: [2]int{scope.id, scope.variablecnt}}
+ return &code{op: opscope, v: [3]int{scope.id, scope.variablecnt, 0}}
  })()
  if c.moduleLoader != nil {
  if moduleLoader, ok := c.moduleLoader.(interface {
@@ -319,7 +319,7 @@ func (c *compiler) compileFuncDef(e *FuncDef, builtin bool) error {
  scope = c.newScope()
  c.scopes = append(c.scopes, scope)
  defer c.lazy(func() *code {
- return &code{op: opscope, v: [2]int{scope.id, scope.variablecnt}}
+ return &code{op: opscope, v: [3]int{scope.id, scope.variablecnt, len(e.Args)}}
  })()
  if len(e.Args) > 0 {
  v := c.newVariable()
@@ -1439,28 +1439,42 @@ func (c *compiler) lazy(f func() *code) func() {
 
 func (c *compiler) optimizeTailRec() {
  var pcs []int
- targets := map[int]struct{}{}
+ scopes := map[int]struct{}{}
+ forked := map[int]struct{}{}
 L:
  for i, l := 0, len(c.codes); i < l; i++ {
  switch c.codes[i].op {
  case opscope:
  pcs = append(pcs, i)
- if c.codes[i].v.([2]int)[1] == 0 {
- targets[i] = struct{}{}
+ if c.codes[i].v.([3]int)[2] == 0 {
+ scopes[i] = struct{}{}
  }
+ case opfork, opforktrybegin, opforkalt:
+ forked[c.codes[i].v.(int)] = struct{}{}
  case opcall:
  if j, ok := c.codes[i].v.(int); !ok ||
  len(pcs) == 0 || pcs[len(pcs)-1] != j {
  break
- } else if _, ok := targets[j]; !ok {
+ } else if _, ok = scopes[j]; !ok {
  break
  }
+ canjump := true
  for j := i + 1; j < l; {
  switch c.codes[j].op {
  case opjump:
  j = c.codes[j].v.(int)
+ if canjump {
+ if _, ok := forked[j+1]; ok {
+ canjump = false
+ }
+ }
  case opret:
- c.codes[i] = &code{op: opjump, v: pcs[len(pcs)-1] + 1}
+ if canjump {
+ c.codes[i].op = opjump
+ c.codes[i].v = pcs[len(pcs)-1] + 1
+ } else {
+ c.codes[i].op = opcallrec
+ }
  continue L
  default:
  continue L

compiler_test.go

@@ -130,7 +130,7 @@ func TestCodeCompile_OptimizeConstants(t *testing.T) {
  }
 }
 
-func TestCodeCompile_OptimizeTailRec(t *testing.T) {
+func TestCodeCompile_OptimizeTailRec_Range(t *testing.T) {
  query, err := gojq.Parse("range(10)")
  if err != nil {
  t.Fatal(err)
@@ -166,6 +166,48 @@ func TestCodeCompile_OptimizeTailRec(t *testing.T) {
  }
 }
 
+func TestCodeCompile_OptimizeTailRec_ScopeVar(t *testing.T) {
+ query, err := gojq.Parse("def f: . as $x | $x, (if $x < 3 then $x + 1 | f else empty end); f")
+ if err != nil {
+ t.Fatal(err)
+ }
+ code, err := gojq.Compile(query)
+ if err != nil {
+ t.Fatal(err)
+ }
+ codes := reflect.ValueOf(code).Elem().FieldByName("codes")
+ if got, expected := codes.Len(), 43; expected != got {
+ t.Errorf("expected: %v, got: %v", expected, got)
+ }
+ op1 := codes.Index(37).Elem().FieldByName("op") // call f by jump
+ op2 := codes.Index(38).Elem().FieldByName("op")
+ if got, expected := *(*int)(unsafe.Pointer(op2.UnsafeAddr())),
+ *(*int)(unsafe.Pointer(op1.UnsafeAddr())); expected != got {
+ t.Errorf("expected: %v, got: %v", expected, got)
+ }
+}
+
+func TestCodeCompile_OptimizeTailRec_CallRec(t *testing.T) {
+ query, err := gojq.Parse("def f: . as $x | $x, (if $x < 3 then $x + 1 | f else empty end), $x; f")
+ if err != nil {
+ t.Fatal(err)
+ }
+ code, err := gojq.Compile(query)
+ if err != nil {
+ t.Fatal(err)
+ }
+ codes := reflect.ValueOf(code).Elem().FieldByName("codes")
+ if got, expected := codes.Len(), 47; expected != got {
+ t.Errorf("expected: %v, got: %v", expected, got)
+ }
+ op1 := codes.Index(38).Elem().FieldByName("op") // callrec f
+ op2 := codes.Index(37).Elem().FieldByName("op") // call _add/2
+ if got, expected := *(*int)(unsafe.Pointer(op2.UnsafeAddr()))+1,
+ *(*int)(unsafe.Pointer(op1.UnsafeAddr())); expected != got {
+ t.Errorf("expected: %v, got: %v", expected, got)
+ }
+}
+
 func TestCodeCompile_OptimizeJumps(t *testing.T) {
  query, err := gojq.Parse("def f: 1; def g: 2; def h: 3; f")
  if err != nil {

debug.go

@@ -75,7 +75,7 @@ func (env *env) debugCodes() {
  for i, c := range env.codes {
  pc := i
  switch c.op {
- case opcall:
+ case opcall, opcallrec:
  if x, ok := c.v.(int); ok {
  pc = x
  }
@@ -87,9 +87,14 @@ func (env *env) debugCodes() {
  }
  var s string
  if name := env.lookupInfoName(pc); name != "" {
- if (c.op == opcall || c.op == opjump) && !strings.HasPrefix(name, "module ") {
- s = "\t## call " + name
- } else {
+ switch c.op {
+ case opcall, opcallrec, opjump:
+ if !strings.HasPrefix(name, "module ") {
+ s = "\t## call " + name
+ break
+ }
+ fallthrough
+ default:
  s = "\t## " + name
  }
  }
@@ -114,7 +119,7 @@ func (env *env) debugState(pc int, backtrack bool) {
  sb.WriteString(debugJSON(env.stack.data[xs[i]].value))
  }
  switch c.op {
- case opcall:
+ case opcall, opcallrec:
  if x, ok := c.v.(int); ok {
  pc = x
  }
@@ -125,9 +130,14 @@ func (env *env) debugState(pc int, backtrack bool) {
  }
  }
  if name := env.lookupInfoName(pc); name != "" {
- if (c.op == opcall || c.op == opjump) && !strings.HasPrefix(name, "module ") {
- sb.WriteString("\t\t\t## call " + name)
- } else {
+ switch c.op {
+ case opcall, opcallrec, opjump:
+ if !strings.HasPrefix(name, "module ") {
+ sb.WriteString("\t\t\t## call " + name)
+ break
+ }
+ fallthrough
+ default:
  sb.WriteString("\t\t\t## " + name)
  }
  }
@@ -162,7 +172,8 @@ func (env *env) debugForks(pc int, op string) {
 }
 
 func debugOperand(c *code) string {
- if c.op == opcall {
+ switch c.op {
+ case opcall, opcallrec:
  switch v := c.v.(type) {
  case int:
  return debugJSON(v)
@@ -171,7 +182,7 @@ func debugOperand(c *code) string {
  default:
  panic(c)
  }
- } else {
+ default:
  return debugJSON(c.v)
  }
 }

execute.go

@@ -187,17 +187,25 @@ loop:
  default:
  panic(v)
  }
+ case opcallrec:
+ pc, callpc, index = code.v.(int), -pc, env.scopes.index
+ goto loop
  case oppushpc:
  env.push([2]int{code.v.(int), env.scopes.index})
  case opcallpc:
  xs := env.pop().([2]int)
  pc, callpc, index = xs[0], pc, xs[1]
  goto loop
  case opscope:
- xs := code.v.([2]int)
+ xs := code.v.([3]int)
  var i, l int
  if index == env.scopes.index {
- i = index
+ if callpc >= 0 {
+ i = index
+ } else {
+ callpc = -callpc
+ i = env.scopes.top().(scope).saveindex
+ }
  } else {
  env.scopes.save(&i, &l)
  env.scopes.index = index