lang: Plumb through a context into unification

If we have a long type unification, we might want to cancel it early.
This also helps us visualize where we want context to be seen.

lang/unification/simplesolver.go

@@ -30,6 +30,7 @@
 package unification // TODO: can we put this solver in a sub-package?
 
 import (
+	"context"
 	"fmt"
 	"sort"
 
@@ -64,9 +65,9 @@ const (
 // SimpleInvariantSolver with the log parameter of your choice specified. The
 // result satisfies the correct signature for the solver parameter of the
 // Unification function.
-func SimpleInvariantSolverLogger(logf func(format string, v ...interface{})) func([]interfaces.Invariant, []interfaces.Expr) (*InvariantSolution, error) {
-	return func(invariants []interfaces.Invariant, expected []interfaces.Expr) (*InvariantSolution, error) {
-		return SimpleInvariantSolver(invariants, expected, logf)
+func SimpleInvariantSolverLogger(logf func(format string, v ...interface{})) func(context.Context, []interfaces.Invariant, []interfaces.Expr) (*InvariantSolution, error) {
+	return func(ctx context.Context, invariants []interfaces.Invariant, expected []interfaces.Expr) (*InvariantSolution, error) {
+		return SimpleInvariantSolver(ctx, invariants, expected, logf)
 	}
 }
 
@@ -195,7 +196,7 @@ func DebugSolverState(solved map[interfaces.Expr]*types.Type, equalities []inter
 
 // SimpleInvariantSolver is an iterative invariant solver for AST expressions.
 // It is intended to be very simple, even if it's computationally inefficient.
-func SimpleInvariantSolver(invariants []interfaces.Invariant, expected []interfaces.Expr, logf func(format string, v ...interface{})) (*InvariantSolution, error) {
+func SimpleInvariantSolver(ctx context.Context, invariants []interfaces.Invariant, expected []interfaces.Expr, logf func(format string, v ...interface{})) (*InvariantSolution, error) {
 	debug := false // XXX: add to interface
 	process := func(invariants []interfaces.Invariant) ([]interfaces.Invariant, []*interfaces.ExclusiveInvariant, error) {
 		equalities := []interfaces.Invariant{}
@@ -395,6 +396,12 @@ func SimpleInvariantSolver(invariants []interfaces.Invariant, expected []interfa
 	// run until we're solved, stop consuming equalities, or type clash
 Loop:
 	for {
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		default:
+			// pass
+		}
 		// Once we're done solving everything else except the generators
 		// then we can exit, but we want to make sure the generators had
 		// a chance to "speak up" and make sure they were part of Unify.
@@ -1314,6 +1321,12 @@ Loop:
 
 			// let's try each combination, one at a time...
 			for i, ex := range exclusivesProduct(exclusives) { // [][]interfaces.Invariant
+				select {
+				case <-ctx.Done():
+					return nil, ctx.Err()
+				default:
+					// pass
+				}
 				logf("%s: exclusive(%d):\n%+v", Name, i, ex)
 				// we could waste a lot of cpu, and start from
 				// the beginning, but instead we could use the
@@ -1324,7 +1337,7 @@ Loop:
 				recursiveInvariants = append(recursiveInvariants, ex...)
 				// FIXME: implement RecursionDepthLimit
 				logf("%s: recursing...", Name)
-				solution, err := SimpleInvariantSolver(recursiveInvariants, expected, logf)
+				solution, err := SimpleInvariantSolver(ctx, recursiveInvariants, expected, logf)
 				if err != nil {
 					logf("%s: recursive solution failed: %+v", Name, err)
 					continue // no solution found here...

lang/unification/simplesolver_test.go

@@ -32,6 +32,7 @@
 package unification
 
 import (
+	"context"
 	"fmt"
 	"strings"
 	"testing"
@@ -265,7 +266,7 @@ func TestSimpleSolver1(t *testing.T) {
 
 			solver := SimpleInvariantSolverLogger(logf) // generates a solver with built-in logging
 
-			solution, err := solver(invariants, expected)
+			solution, err := solver(context.TODO(), invariants, expected)
 			t.Logf("test #%d: solver completed with: %+v", index, err)
 
 			if !fail && err != nil {

lang/unification/unification.go

@@ -32,6 +32,7 @@
 package unification
 
 import (
+	"context"
 	"fmt"
 	"sort"
 	"strings"
@@ -45,7 +46,8 @@ type Unifier struct {
 	AST interfaces.Stmt
 
 	// Solver is the solver algorithm implementation to use.
-	Solver func([]interfaces.Invariant, []interfaces.Expr) (*InvariantSolution, error)
+	// XXX: Solver should be a solver interface, not a function signature.
+	Solver func(context.Context, []interfaces.Invariant, []interfaces.Expr) (*InvariantSolution, error)
 
 	Debug bool
 	Logf  func(format string, v ...interface{})
@@ -63,7 +65,7 @@ type Unifier struct {
 // type. This function and logic was invented after the author could not find
 // any proper literature or examples describing a well-known implementation of
 // this process. Improvements and polite recommendations are welcome.
-func (obj *Unifier) Unify() error {
+func (obj *Unifier) Unify(ctx context.Context) error {
 	if obj.AST == nil {
 		return fmt.Errorf("the AST is nil")
 	}
@@ -96,7 +98,7 @@ func (obj *Unifier) Unify() error {
 	exprMap := ExprListToExprMap(exprs)    // makes searching faster
 	exprList := ExprMapToExprList(exprMap) // makes it unique (no duplicates)
 
-	solved, err := obj.Solver(invariants, exprList)
+	solved, err := obj.Solver(ctx, invariants, exprList)
 	if err != nil {
 		return err
 	}