internal/unify: new package for structured value unification

The plan is to use this package as part of generating SIMD mappings.

Change-Id: Ie67bf7fe87222b8dffdbb12a99729c0fe0f7bc38
Reviewed-on: https://go-review.googlesource.com/c/arch/+/666515
LUCI-TryBot-Result: Go LUCI <[email protected]>
Reviewed-by: Junyang Shao <[email protected]>

Author: aclements

go.mod

@@ -3,3 +3,5 @@ module golang.org/x/arch
 go 1.23.0
 
 require rsc.io/pdf v0.1.1
+
+require gopkg.in/yaml.v3 v3.0.1

go.sum

@@ -1,2 +1,6 @@
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 rsc.io/pdf v0.1.1 h1:k1MczvYDUvJBe93bYd7wrZLLUEcLZAuF824/I4e5Xr4=
 rsc.io/pdf v0.1.1/go.mod h1:n8OzWcQ6Sp37PL01nO98y4iUCRdTGarVfzxY20ICaU4=

internal/unify/closure.go

@@ -0,0 +1,148 @@
+// Copyright 2025 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package unify
+
+import (
+	"fmt"
+	"iter"
+	"maps"
+	"slices"
+)
+
+type Closure struct {
+	val *Value
+	env nonDetEnv
+}
+
+func NewSum(vs ...*Value) Closure {
+	id := &ident{name: "sum"}
+	return Closure{NewValue(Var{id}), topEnv.bind(id, vs...)}
+}
+
+// IsBottom returns whether c consists of no values.
+func (c Closure) IsBottom() bool {
+	return c.val.Domain == nil
+}
+
+// Summands returns the top-level Values of c. This assumes the top-level of c
+// was constructed as a sum, and is mostly useful for debugging.
+func (c Closure) Summands() iter.Seq[*Value] {
+	if v, ok := c.val.Domain.(Var); ok {
+		parts := c.env.partitionBy(v.id)
+		return func(yield func(*Value) bool) {
+			for _, part := range parts {
+				if !yield(part.value) {
+					return
+				}
+			}
+		}
+	}
+	return func(yield func(*Value) bool) {
+		yield(c.val)
+	}
+}
+
+// All enumerates all possible concrete values of c by substituting variables
+// from the environment.
+//
+// E.g., enumerating this Value
+//
+//	a: !sum [1, 2]
+//	b: !sum [3, 4]
+//
+// results in
+//
+//   - {a: 1, b: 3}
+//   - {a: 1, b: 4}
+//   - {a: 2, b: 3}
+//   - {a: 2, b: 4}
+func (c Closure) All() iter.Seq[*Value] {
+	// In order to enumerate all concrete values under all possible variable
+	// bindings, we use a "non-deterministic continuation passing style" to
+	// implement this. We use CPS to traverse the Value tree, threading the
+	// (possibly narrowing) environment through that CPS following an Euler
+	// tour. Where the environment permits multiple choices, we invoke the same
+	// continuation for each choice. Similar to a yield function, the
+	// continuation can return false to stop the non-deterministic walk.
+	return func(yield func(*Value) bool) {
+		c.val.all1(c.env, func(v *Value, e nonDetEnv) bool {
+			return yield(v)
+		})
+	}
+}
+
+func (v *Value) all1(e nonDetEnv, cont func(*Value, nonDetEnv) bool) bool {
+	switch d := v.Domain.(type) {
+	default:
+		panic(fmt.Sprintf("unknown domain type %T", d))
+
+	case nil:
+		return true
+
+	case Top, String:
+		return cont(v, e)
+
+	case Def:
+		fields := d.keys()
+		// We can reuse this parts slice because we're doing a DFS through the
+		// state space. (Otherwise, we'd have to do some messy threading of an
+		// immutable slice-like value through allElt.)
+		parts := make(map[string]*Value, len(fields))
+
+		// TODO: If there are no Vars or Sums under this Def, then nothing can
+		// change the Value or env, so we could just cont(v, e).
+		var allElt func(elt int, e nonDetEnv) bool
+		allElt = func(elt int, e nonDetEnv) bool {
+			if elt == len(fields) {
+				// Build a new Def from the concrete parts. Clone parts because
+				// we may reuse it on other non-deterministic branches.
+				nVal := newValueFrom(Def{maps.Clone(parts)}, v)
+				return cont(nVal, e)
+			}
+
+			return d.fields[fields[elt]].all1(e, func(v *Value, e nonDetEnv) bool {
+				parts[fields[elt]] = v
+				return allElt(elt+1, e)
+			})
+		}
+		return allElt(0, e)
+
+	case Tuple:
+		// Essentially the same as Def.
+		if d.repeat != nil {
+			// There's nothing we can do with this.
+			return cont(v, e)
+		}
+		parts := make([]*Value, len(d.vs))
+		var allElt func(elt int, e nonDetEnv) bool
+		allElt = func(elt int, e nonDetEnv) bool {
+			if elt == len(d.vs) {
+				// Build a new tuple from the concrete parts. Clone parts because
+				// we may reuse it on other non-deterministic branches.
+				nVal := newValueFrom(Tuple{vs: slices.Clone(parts)}, v)
+				return cont(nVal, e)
+			}
+
+			return d.vs[elt].all1(e, func(v *Value, e nonDetEnv) bool {
+				parts[elt] = v
+				return allElt(elt+1, e)
+			})
+		}
+		return allElt(0, e)
+
+	case Var:
+		// Go each way this variable can be bound.
+		for _, ePart := range e.partitionBy(d.id) {
+			// d.id is no longer bound in this environment partition. We'll may
+			// need it later in the Euler tour, so bind it back to this single
+			// value.
+			env := ePart.env.bind(d.id, ePart.value)
+			if !ePart.value.all1(env, cont) {
+				return false
+			}
+		}
+		return true
+	}
+}