Implement tail calls.

Author: Molorius

pkg/forth/basic_test.go

@@ -164,6 +164,36 @@ func TestPrimitives(t *testing.T) {
 			asm:    wrapMain("4 0 DO I . LOOP"),
 			expect: "0 1 2 3 ",
 		},
+		{
+			name: "tail call",
+			asm: `
+				\ The point of this isn't to test
+				\ a forth algorithm, it's to test that the
+				\ compiler is succesfully able to generate
+				\ a tail call.
+				\ If it doesn't, the return stack pointer
+				\ will overflow on the ulp.
+
+				\ countdown from n to 0
+				: countdown ( n -- )
+					DUP 0= IF
+						DROP \ we don't care about the result
+					ELSE
+						1-
+						RECURSE \ tail call recursively!
+						EXIT \ TODO get this to work without the EXIT
+					THEN
+				;
+
+				: MAIN
+					DEPTH u. \ print the depth
+					5 countdown DEPTH u. \ try after a few recursive calls
+					2000 countdown DEPTH u. \ crank it up to 11
+					ESP.DONE
+				;
+			`,
+			expect: "0 0 0 ",
+		},
 	}
 	r := asm.Runner{}
 	r.SetDefaults()

pkg/forth/cell.go

@@ -332,3 +332,42 @@ func (c *CellBranch0) IsRecursive(check *WordForth) bool {
 func (c *CellBranch0) String() string {
 	return fmt.Sprintf("Branch0{%p}", c.dest)
 }
+
+// Used during an optimization pass to add
+// tail calls.
+type CellTailCall struct {
+	dest *WordForth
+}
+
+func (c *CellTailCall) Execute(vm *VirtualMachine) error {
+	return fmt.Errorf("Cannot directly execute a tail call, please file a bug repot")
+}
+
+func (c *CellTailCall) AddToList(u *Ulp) error {
+	// Add the destination to the list.
+	return c.dest.AddToList(u)
+}
+
+func (c *CellTailCall) BuildExecution(u *Ulp) (string, error) {
+	switch u.compileTarget {
+	case UlpCompileTargetToken:
+		return fmt.Sprintf(".int %s + 0x8000", c.dest.Entry.ulpName), nil
+	case UlpCompileTargetSubroutine:
+		// put the address after the docol
+		return fmt.Sprintf("move r2, %s + 1\r\njump r2", c.dest.Entry.ulpName), nil
+	default:
+		return "", fmt.Errorf("Unknown compile target %d, please file a bug report", u.compileTarget)
+	}
+}
+
+func (c *CellTailCall) OutputReference(u *Ulp) (string, error) {
+	return "", fmt.Errorf("Cannot refer to a tail call, please file a bug report")
+}
+
+func (c *CellTailCall) IsRecursive(check *WordForth) bool {
+	return c.dest.IsRecursive(check)
+}
+
+func (c *CellTailCall) String() string {
+	return fmt.Sprintf("TailCall{%s}", c.dest.Entry.Name)
+}

pkg/forth/flag.go

@@ -16,4 +16,5 @@ type Flag struct {
 	addedToList bool // This word is already added to the output list.
 	recursive   bool // This Forth word is recursive.
 	visited     bool // This Forth word has already been visited in this optimization pass.
+	isExit      bool // This assembly word is the EXIT word.
 }

pkg/forth/primitive.go

@@ -1058,6 +1058,9 @@ func PrimitiveSetup(vm *VirtualMachine) error {
 		},
 		{
 			name: "EXIT",
+			flag: Flag{
+				isExit: true,
+			},
 			goFunc: func(vm *VirtualMachine, entry *DictionaryEntry) error {
 				ret, err := vm.ReturnStack.Pop()
 				if err != nil {

pkg/forth/ulpBuild.go

@@ -207,6 +207,10 @@ func (u *Ulp) optimize() error {
 	if err != nil {
 		return errors.Join(fmt.Errorf("could not tag recursion during optimization"), err)
 	}
+	err = u.putTailCalls()
+	if err != nil {
+		return errors.Join(fmt.Errorf("could not create tail calls, please file a bug report"), err)
+	}
 	return nil
 }
 
@@ -228,6 +232,40 @@ func (u *Ulp) tagRecursion() error {
 	return nil
 }
 
+// Put tail calls as an optimization.
+func (u *Ulp) putTailCalls() error {
+	for _, w := range u.forthWords {
+		length := len(w.Cells) - 1
+		for i := 0; i < length; i++ {
+			// check if the first cell is a forth word
+			firstAddress, ok := w.Cells[i].(CellAddress)
+			if !ok {
+				continue
+			}
+			word, ok := firstAddress.Entry.Word.(*WordForth)
+			if !ok {
+				continue
+			}
+			// check if the second cell is an EXIT
+			secondAddress, ok := w.Cells[i+1].(CellAddress)
+			if !ok {
+				continue
+			}
+			if !secondAddress.Entry.Flag.isExit {
+				continue
+			}
+			// replace both cells with the tail call!
+			tailCall := CellTailCall{word}     // create the tail call
+			w.Cells[i] = &tailCall             // replace the word
+			before := w.Cells[:i+1]            // get the cells before, including the tail call
+			after := w.Cells[i+2:]             // get the cells after, excluding the exit
+			w.Cells = append(before, after...) // recreate the list
+			length -= 1                        // shift the length
+		}
+	}
+	return nil
+}
+
 func (u *Ulp) clearVisited() {
 	for _, w := range u.forthWords {
 		w.Entry.ClearVisited()