expr.go

package shapes

import (
	"fmt"
)

//go-sumtype:decl Expr

// Expr represents an expression. The following types are Expr:
//
//	Shape | Abstract | Arrow | Compound
//	Var | Size | UnaryOp
//	IndexOf | TransposeOf | SliceOf | RepeatOf | ConcatOf
//	Sli | Axis | Axes
//
// A compact BNF is as follows:
//
//	E := S | A | E → E | (E s.t. X)
//	a | Sz | Π E | Σ E | D E
//	I n E | T []Ax E | L : E | R Ax n E | C Ax E E
//	: | Ax | []Ax
type Expr interface {
	isExpr()
	depth() int // depth of the expression at at that point

	substitutable
}

// Var represents a variable. A variable can represent:
//   - a Shape (e.g. a in the expression a → b)
//   - a Size (e.g. (a, b))
//   - a Slice (in Arrow expressions)
type Var rune

func (v Var) isSizelike()                {}
func (v Var) isSlicelike()               {}
func (v Var) isExpr()                    {}
func (v Var) Format(s fmt.State, r rune) { fmt.Fprintf(s, "%c", rune(v)) }
func (v Var) apply(ss substitutions) substitutable {
	if len(ss) == 0 {
		return v
	}
	for _, s := range ss {
		if s.For == v {
			return s.Sub
		}
	}
	return v
}
func (v Var) freevars() varset              { return varset{v} }
func (v Var) subExprs() []substitutableExpr { return nil }
func (v Var) depth() int                    { return 0 }

// Axis represents an axis in doing shape stuff.
type Axis int

const (
	AllAxes Axis = -65535
)

func (a Axis) isExpr()                              {}
func (a Axis) apply(ss substitutions) substitutable { return a }
func (a Axis) freevars() varset                     { return nil }
func (a Axis) subExprs() []substitutableExpr        { return nil }
func (a Axis) depth() int                           { return 0 }
func (a Axis) Format(s fmt.State, c rune) {
	if c == 'x' {
		if a == AllAxes {
			fmt.Fprintf(s, "⁼")
			return
		}
		fmt.Fprintf(s, supInt(int(a)))
		return
	}
	if a == AllAxes {
		fmt.Fprintf(s, ":")
		return
	}
	fmt.Fprintf(s, "%d", int(a))
}

func ResolveAxis(a Axis, s Shapelike) Axis {
	if a == AllAxes {
		return AllAxes
	}
	if a < 0 {
		return Axis(s.Dims() + int(a))
	}
	return a
}

// Axes represents a list of axes.
// Despite being a container type (i.e. an Axis is an Expr),
// it returns nil for Exprs(). This is because we want to treat Axes as a monolithic entity.
type Axes []Axis

func (a Axes) isExpr() {}
func (a Axes) Format(s fmt.State, r rune) {
	if r == 'x' {
		fmt.Fprintf(s, supInts(axesToInts(a)))
		return
	}
	fmt.Fprintf(s, "X%v", axesToInts(a))
}
func (a Axes) apply(ss substitutions) substitutable { return a }
func (a Axes) freevars() varset                     { return nil }
func (a Axes) subExprs() []substitutableExpr        { return nil }
func (a Axes) depth() int                           { return 1 } // Axes is a single term
func (a Axes) Dims() int                            { return len(a) }
func (a Axes) AsInts() []int                        { return axesToInts(a) }
func (a Axes) Eq(other Axes) bool {
	return len(a) == len(other) && intsEq(axesToInts(a), axesToInts(other))
}

// Size represents a size of a dimension/axis
type Size int

func (s Size) isExpr()                              {}
func (s Size) isSizelike()                          {}
func (s Size) apply(ss substitutions) substitutable { return s }
func (s Size) freevars() varset                     { return nil }
func (s Size) subExprs() []substitutableExpr        { return nil }
func (s Size) depth() int                           { return 1 }
func (s Size) Format(f fmt.State, r rune)           { fmt.Fprintf(f, "%d", int(s)) }

// Size also implements Operation (i.e. it's a Const)

func (s Size) isValid() bool              { return true }
func (s Size) resolveSize() (Size, error) { return s, nil }

// Sizes are a list of sizes.
type Sizes []Size

func (s Sizes) isExpr()                              {}
func (s Sizes) Format(f fmt.State, r rune)           { fmt.Fprintf(f, "Sz%v", sizesToInts(s)) }
func (s Sizes) apply(ss substitutions) substitutable { return s }
func (s Sizes) freevars() varset                     { return nil }
func (s Sizes) subExprs() []substitutableExpr        { return nil }
func (s Sizes) depth() int                           { return 1 }
func (s Sizes) AsInts() []int                        { return sizesToInts(s) }

// complex expressions

// Arrow represents a function of shapes, from A → B. Arrows are right associative.
type Arrow struct {
	A, B Expr
}

// MakeArrow is a utility function for writing correct Arrow expressions.
//
// Consider for example, matrix multiplication. It is written plainly as follows:
//
//	MatMul: (a, b) → (b, c) → (a, c)
//
// However, note that because arrows are right associative and they're a binary operator, it
// actually is more correctly written like this:
//
//	MatMul: (a, b) → ((b, c) → (a, c))
//
// This makes writing plain Arrow{} expressions a bit fraught with errors (see example).
// Thus, the MakeArrow function is created to help write more correct expressions.
func MakeArrow(exprs ...Expr) Arrow {
	if len(exprs) < 2 {
		panic("Expect at least two expressions to make an Arrow")
	}
	a := Arrow{A: exprs[0]}
	if len(exprs) > 2 {
		a.B = MakeArrow(exprs[1:]...)
	} else {
		a.B = exprs[1]
	}
	return a
}

func (a Arrow) isExpr() {}

func (a Arrow) depth() int { return max(a.A.depth(), a.B.depth()) + 1 }

func (a Arrow) Format(s fmt.State, r rune) {
	if _, ok := a.A.(Arrow); ok {
		fmt.Fprintf(s, "(%v) → %v", a.A, a.B)
		return
	}
	fmt.Fprintf(s, "%v → %v", a.A, a.B)
}

func (a Arrow) apply(ss substitutions) substitutable {
	return Arrow{
		A: a.A.apply(ss).(Expr),
		B: a.B.apply(ss).(Expr),
	}
}
func (a Arrow) freevars() varset { return arrowToTup(&a).freevars() }

func (a Arrow) subExprs() []substitutableExpr {
	return []substitutableExpr{a.A.(substitutableExpr), a.B.(substitutableExpr)}
}

/* Example

MatMul:
(a, b) → (b, c) → (a, c)

is represented as:

Arrow {
	Arrow {
		Abstract{Var('a'), Var('b')},
		Abstract{Var('b'), Var('c')},
	},
	Abstract{Var('a'), Var('c')},
}

Add:

a → a → a

is represented as:

Arrow {
	Arrow {
		Var('a'),
		Var('a'),
	},
	Var('a'),
}


Flatten/Ravel:
a → Πa

is represented as:

Arrow {
	Var('a'),
	UnaryOp {Prod, Var('a')},
}


Sum:

a → ()

is represented as:

Arrow {
	Var('a'),
	Shape{},
}

Transpose:
a → Axes → Tr Axes a

is represented as:

Arrow {
	Arrow{
		Var('a'),
		axes,
	},
	TransposeOf {
		axes,
		Var('a'),
	},
}

Slice:
a → Sli → SliceOf Sli a

if represented as:

Arrow {
	Arrow {
		Var('a'),
		sli,
	},
	SliceOf {
		sli,
		Var('a'),
	}
}


More complicated examples:

Reshape:

a → b → b s.t. (Πa = Πb)

is represented as

Compound {
	Arrow {
		Arrow {
			Var('a'),
			Var('b'),
		},
		Var('b'),
	},
	SubjectTo {
		Eq,
		UnaryOp{Prod, Var('a')},
		UnaryOp{Prod, Var('b')},
	}
}


Sum a matrix columnwise (along axis 1):

(a → Axes→ Tr Axes a) → b s.t. (D b = D a - 1)

is represented as:

Compound {
	Arrow {
		Arrow {
			Arrow {
				Var('a'),
				Axes{1,0},
			},
			TransposeOf {
				Axes{1, 0},
				Var('a')
			},
		},
		Var('b'),
	},
	SubjectTo {
		Eq,
		UnaryOp{Dim, Var('b')},
		BinaryOp {
			Sub,
			UnaryOp {Dim, Var('a')},
			Size(1),
		}
	},
}

*/