make IF/ELSE/THEN work!

builtins.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"errors"
 	"fmt"
 	"io"
 	"os"
@@ -23,7 +24,7 @@ func (b *Builtins) Colon(c *Context) func() error {
 // Semicolon sets the COMPILE/IMMEDIATE flag back to IMMEDIATE and adds the defintion to the dictionary
 func (b *Builtins) Semicolon(c *Context) func() error {
 	return func() error {
-		c.Dictionary.AddWord(c.Words[0], nil, c.Words[1:], false)
+		c.Dictionary.AddWord(c.Words[0], Word{Name: c.Words[0], Source: c.Words[1:]})
 		c.Words = []string{}
 		c.Flags.SetFlag("Immediate", true)
 		return nil
@@ -391,24 +392,62 @@ func (b *Builtins) RFetch(s *Stack, r *Stack) func() error {
 	}
 }
 
-// If checks if NOS != 0 and conditionally executes any following statements if so
+// If checks if TOS != 0 and executes any following statements if it is. If TOS == 0, skip execution until ELSE or THEN.
-func (b *Builtins) If(s *Stack, r *Stack) func() error {
+// Nested IFs and ELSE/THENs inside "non-executing" blocks still manipulate the if stack, in order to properly balance
+// IFs and THENs. i.e. `0 IF 0 IF 42 EMIT THEN 42 EMIT THEN` should not print anything, but a naive "go until you see any THEN"
+// would output "* ".
+func (b *Builtins) If(s *Stack, i *Stack) func() error {
 	return func() error {
+		// check to see the current IF state. If we're inside a non-executing branch,
+		// add our own non-execution state, AND tell ELSE to skip its work to boot.
+		itop, err := i.Pick(0)
+		if err != nil && !errors.Is(err, ErrUnderflow) {
+			return err
+		} else if err == nil && itop == 0 {
+			i.Push(-1)
+			return nil
+		}
+
+		tos, err := s.Pop()
+		if err != nil {
+			return err
+		}
+		// TOS was 0, don't execute this branch
+		if tos == 0 {
+			i.Push(0)
+		} else {
+			// TOS wasn't 0, execute until we see ELSE/THEN
+			i.Push(1)
+		}
 		return nil
 	}
 }
 
-// Else executes a separate piece of code if NOS == 0 after an if check
+// Else executes a separate piece of code if the IF stack
-func (b *Builtins) Else(s *Stack, r *Stack) func() error {
+func (b *Builtins) Else(i *Stack) func() error {
 	return func() error {
+		itop, err := i.Pop()
+		if err != nil {
+			return err
+		}
+		// don't execute until we see our THEN
+		if itop == -1 || itop == 1 {
+			i.Push(0)
+		}
+		// we weren't running code, but now we can until THEN
+		if itop == 0 {
+			i.Push(1)
+		}
 		return nil
 	}
 }
 
 // Then ends an If/Else block
-func (b *Builtins) Then(s *Stack, r *Stack) func() error {
+func (b *Builtins) Then(i *Stack) func() error {
 	return func() error {
-		return nil
+		// Pop off the existing IF/ELSE state and return
+		_, err := i.Pop()
+		return err
 	}
 }
 

eval.go

@@ -9,8 +9,9 @@ import (
 // Context is a set of Dictionary + Stacks + Flags representing a runtime environment
 type Context struct {
 	Dictionary Dictionary
-	Stack      *Stack
+	Stack      *Stack // main stack
-	RStack     *Stack
+	RStack     *Stack // "return" stack, for handling loops/backward jumps
+	IfStack    *Stack // tracks nested branches
 	Flags      Flags
 	Words      []string
 }
@@ -43,27 +44,25 @@ func (c *Context) Eval(line string) error {
 				continue
 			}
 
-			int, err := strconv.Atoi(sword)
+			ifcheck, _ := c.IfStack.Pick(0)
-			if err == nil {
+			if len(c.IfStack.values) == 0 || (len(c.IfStack.values) > 0 && ifcheck == 1) || w.BranchCheck {
-				// it was a number! put it on the stack.
+				int, err := strconv.Atoi(sword)
-				c.Stack.Push(int)
+				if err == nil {
-				word = []byte{}
+					// it was a number! put it on the stack.
-				continue
+					c.Stack.Push(int)
-			}
+					word = []byte{}
+					continue
+				}
 
-			// it wasn't a number. Is it a word we know?
+				// run word
-			w, err = c.Dictionary.GetWord(sword)
+				c.RStack.Push(i)
-			if err != nil {
+				if err = c.Exec(w); err != nil {
-				return fmt.Errorf("could not parse %s; %v", w.Name, err)
+					return err
-			}
+				}
-			// run word
+				i, err = c.RStack.Pop()
-			c.RStack.Push(i)
+				if err != nil {
-			if err = c.Exec(w); err != nil {
+					return fmt.Errorf("error while popping from return stack: %v", err)
-				return err
+				}
-			}
-			i, err = c.RStack.Pop()
-			if err != nil {
-				return fmt.Errorf("error while popping from return stack: %v", err)
 			}
 			word = []byte{}
 		default:

main.go

@@ -11,6 +11,7 @@ import (
 func main() {
 	stack := Stack{values: []int{}}
 	rstack := Stack{values: []int{}}
+	ifstack := Stack{values: []int{}}
 
 	dict := Dictionary{}
 
@@ -18,6 +19,7 @@ func main() {
 		Dictionary: dict,
 		Stack:      &stack,
 		RStack:     &rstack,
+		IfStack:    &ifstack,
 		Flags: Flags{
 			"Immediate": true,
 			"Comment":   false,
@@ -27,50 +29,54 @@ func main() {
 
 	b := &Builtins{}
 	// word definitions
-	dict.AddWord(":", b.Colon(&c), nil, false)
+	dict.AddWord(":", Word{Name: ":", Impl: b.Colon(&c)})
-	dict.AddWord(";", b.Semicolon(&c), nil, true)
+	dict.AddWord(";", Word{Name: ";", Impl: b.Semicolon(&c), Immediate: true})
 	// comments
-	dict.AddWord("(", b.OpenComment(&c), nil, true)
+	dict.AddWord("(", Word{Name: "(", Impl: b.OpenComment(&c), Immediate: true})
-	dict.AddWord(")", b.CloseComment(&c), nil, true)
+	dict.AddWord(")", Word{Name: ")", Impl: b.CloseComment(&c), Immediate: true})
 	// math
-	dict.AddWord("+", b.Add(&stack), nil, false)
+	dict.AddWord("+", Word{Name: "+", Impl: b.Add(&stack)})
-	dict.AddWord("-", b.Sub(&stack), nil, false)
+	dict.AddWord("-", Word{Name: "-", Impl: b.Sub(&stack)})
-	dict.AddWord("*", b.Mul(&stack), nil, false)
+	dict.AddWord("*", Word{Name: "*", Impl: b.Mul(&stack)})
-	dict.AddWord("/", b.Div(&stack), nil, false)
+	dict.AddWord("/", Word{Name: "/", Impl: b.Div(&stack)})
 	// output
-	dict.AddWord(".", b.Print(os.Stdout, &stack), nil, false)
+	dict.AddWord(".", Word{Name: ".", Impl: b.Print(os.Stdout, &stack)})
-	dict.AddWord("EMIT", b.Emit(os.Stdout, &stack), nil, false)
+	dict.AddWord("EMIT", Word{Name: "EMIT", Impl: b.Emit(os.Stdout, &stack)})
-	dict.AddWord("CR", nil, []string{"10", "EMIT"}, false) // emit a newline
+	dict.AddWord("CR", Word{Name: "CR", Source: []string{"10", "EMIT"}}) // emit a newline
 	// logic
-	dict.AddWord("=", b.Eq(&stack), nil, false)
+	dict.AddWord("=", Word{Name: "=", Impl: b.Eq(&stack)})
-	dict.AddWord("0=", nil, []string{"0", "="}, false)
+	dict.AddWord("0=", Word{Name: "0=", Source: []string{"0", "="}})
-	dict.AddWord("<>", b.NEq(&stack), nil, false)
+	dict.AddWord("<>", Word{Name: "<>", Impl: b.NEq(&stack)})
-	dict.AddWord(">", b.Gt(&stack), nil, false)
+	dict.AddWord(">", Word{Name: ">", Impl: b.Gt(&stack)})
-	dict.AddWord("<", b.Lt(&stack), nil, false)
+	dict.AddWord("<", Word{Name: "<", Impl: b.Lt(&stack)})
-	dict.AddWord(">=", b.GtEq(&stack), nil, false)
+	dict.AddWord(">=", Word{Name: ">=", Impl: b.GtEq(&stack)})
-	dict.AddWord("<=", b.LtEq(&stack), nil, false)
+	dict.AddWord("<=", Word{Name: "<=", Impl: b.LtEq(&stack)})
-	dict.AddWord("0<", nil, []string{"0", "<"}, false)
+	dict.AddWord("0<", Word{Name: "0<", Source: []string{"0", "<"}})
-	dict.AddWord("0>", nil, []string{"0", ">"}, false)
+	dict.AddWord("0>", Word{Name: "0>", Source: []string{"0", ">"}})
 	// stack manipulation
-	dict.AddWord("DUP", b.Dup(&stack), nil, false)
+	dict.AddWord("DUP", Word{Name: "DUP", Impl: b.Dup(&stack)})
-	dict.AddWord("SWAP", b.Swap(&stack), nil, false)
+	dict.AddWord("SWAP", Word{Name: "SWAP", Impl: b.Swap(&stack)})
-	dict.AddWord("OVER", b.Over(&stack), nil, false)
+	dict.AddWord("OVER", Word{Name: "OVER", Impl: b.Over(&stack)})
-	dict.AddWord("DROP", b.Drop(&stack), nil, false)
+	dict.AddWord("DROP", Word{Name: "DROP", Impl: b.Drop(&stack)})
-	dict.AddWord("ROT", b.Rot(&stack), nil, false)
+	dict.AddWord("ROT", Word{Name: "ROT", Impl: b.Rot(&stack)})
 	// debugging
-	dict.AddWord("WORDS", b.Words(dict), nil, false)
+	dict.AddWord("WORDS", Word{Name: "WORDS", Impl: b.Words(dict)})
-	dict.AddWord("FLAGS", b.Flags(c), nil, false)
+	dict.AddWord("FLAGS", Word{Name: "FLAGS", Impl: b.Flags(c)})
-	dict.AddWord(".S", b.Debug(&stack), nil, false)
+	dict.AddWord(".S", Word{Name: ".S", Impl: b.Debug(&stack)})
-	dict.AddWord(".R", b.Debug(&rstack), nil, false)
+	dict.AddWord(".R", Word{Name: ".R", Impl: b.Debug(&rstack)})
-	dict.AddWord("R>", b.RFrom(&stack, &rstack), nil, false)
+	dict.AddWord(".I", Word{Name: ".I", Impl: b.Debug(&ifstack)})
-	dict.AddWord(">R", b.ToR(&stack, &rstack), nil, false)
+	dict.AddWord("R>", Word{Name: "R>", Impl: b.RFrom(&stack, &rstack)})
-	dict.AddWord("R@", b.RFetch(&stack, &rstack), nil, false)
+	dict.AddWord(">R", Word{Name: ">R", Impl: b.ToR(&stack, &rstack)})
+	dict.AddWord("R@", Word{Name: "R@", Impl: b.RFetch(&stack, &rstack)})
 	// branching
-	dict.AddWord("DO", b.Do(&stack, &rstack), nil, false)
+	dict.AddWord("IF", Word{Name: "IF", Impl: b.If(&stack, &ifstack), BranchCheck: true})
-	dict.AddWord("LOOP", b.Loop(&stack, &rstack), nil, false)
+	dict.AddWord("ELSE", Word{Name: "ELSE", Impl: b.Else(&ifstack), BranchCheck: true})
-	dict.AddWord("I", b.I(&stack, &rstack), nil, false)
+	dict.AddWord("THEN", Word{Name: "THEN", Impl: b.Then(&ifstack), BranchCheck: true})
+	dict.AddWord("DO", Word{Name: "DO", Impl: b.Do(&stack, &rstack)})
+	dict.AddWord("LOOP", Word{Name: "LOOP", Impl: b.Loop(&stack, &rstack)})
+	dict.AddWord("I", Word{Name: "I", Impl: b.I(&stack, &rstack)})
 	// exit
-	dict.AddWord("BYE", b.Quit(), nil, false)
+	dict.AddWord("BYE", Word{Name: "BYE", Impl: b.Quit()})
 
 	reader := bufio.NewReader(os.Stdin)
 	fmt.Print("prosper\n")

stack.go

@@ -2,6 +2,12 @@ package main
 
 import "fmt"
 
+// ErrEmpty denotes an empty stack
+var ErrEmpty = fmt.Errorf("stack empty")
+
+// ErrUnderflow denotes a PICK that tried to pick a value that didn't exist
+var ErrUnderflow = fmt.Errorf("cannot pick value from beyond stack depth")
+
 // Stack is a stack of integers with no defined max depth
 type Stack struct {
 	values []int
@@ -10,7 +16,7 @@ type Stack struct {
 // Pop returns the top of the stack
 func (s *Stack) Pop() (int, error) {
 	if len(s.values) == 0 {
-		return 0, fmt.Errorf("stack empty")
+		return 0, ErrEmpty
 	}
 	i := s.values[0]
 	s.values = s.values[1:]
@@ -24,8 +30,8 @@ func (s *Stack) Push(i int) {
 
 // Pick grabs a value from the given depth in the stack, non-destructively.
 func (s *Stack) Pick(i int) (int, error) {
-	if len(s.values) < i {
+	if len(s.values) <= i {
-		return 0, fmt.Errorf("cannot pick value from beyond stack depth")
+		return 0, ErrUnderflow
 	}
 	return s.values[i], nil
 }

stack_test.go

@@ -42,3 +42,25 @@ func TestStackPush(t *testing.T) {
 		t.Fail()
 	}
 }
+
+func TestStackPick(t *testing.T) {
+	s := Stack{}
+	i, err := s.Pick(0)
+	if err == nil {
+		t.Fail()
+	}
+	if i != 0 {
+		t.Fail()
+	}
+
+	s = Stack{
+		values: []int{1, 2, 3},
+	}
+	i, err = s.Pick(2)
+	if err != nil {
+		t.Fail()
+	}
+	if i != 3 {
+		t.Fail()
+	}
+}

words.go

@@ -2,25 +2,24 @@ package main
 
 import "fmt"
 
+// ErrNoWordFound flags if the requested word is not in the dictionary
+var ErrNoWordFound = fmt.Errorf("no word found")
+
 // Dictionary is a simple map of names to words
 type Dictionary map[string]Word
 
 // A Word defines a subroutine
 type Word struct {
-	Name      string
+	Name        string
-	Impl      func() error
+	Impl        func() error
-	Source    []string
+	Source      []string
-	Immediate bool
+	Immediate   bool
+	BranchCheck bool
 }
 
-// AddWord inserts a new word into the dictonary, optionally overwriting the existing word
+// AddWord inserts a new word into the dictonary, overwriting any existing word by that name
-func (d Dictionary) AddWord(name string, impl func() error, source []string, immediate bool) {
+func (d Dictionary) AddWord(name string, w Word) {
-	d[name] = Word{
+	d[name] = w
-		Name:      name,
-		Impl:      impl,
-		Source:    source,
-		Immediate: immediate,
-	}
 }
 
 // GetWord returns a word from the dictionary or an error if it's undefined
@@ -29,7 +28,7 @@ func (d Dictionary) GetWord(name string) (Word, error) {
 	if !ok {
 		return Word{
 			Name: name,
-		}, fmt.Errorf("no word found")
+		}, ErrNoWordFound
 	}
 	return w, nil
 }

words_test.go

@@ -5,7 +5,10 @@ import "testing"
 func TestDictionary(t *testing.T) {
 	d := Dictionary{}
 
-	d.AddWord("INC", nil, []string{"1", "+"}, false)
+	d.AddWord("INC", Word{
+		Name:   "INC",
+		Source: []string{"1", "+"},
+	})
 	w, err := d.GetWord("INC")
 	if err != nil {
 		t.Fail()