JSON-Lang/ir.py

from emis_funky_funktions import *

from typing import Mapping, Sequence, Tuple, TypeAlias


Expression: TypeAlias = 'Function | Application | Int | Variable | Builtin'
Pattern: TypeAlias = 'NamePattern | IntPattern | SPattern | IgnorePattern'

@dataclass(frozen=True)
class NamePattern:
	"""
	A pattern which always succeeds to match, and binds a whole expression to a name
	"""
	name: str

	def binds(self, var: str) -> bool:
		"""
		Test to see if this pattern binds a given variable
		"""
		return var == self.name

	def match(self, e: Expression) -> Option[Sequence[Tuple[str, Expression]]]:
		"""
		Match an expression against this pattern

		>>> NamePattern('my_var').match(Int(1))
		Some((('my_var', 1),))
		"""
		return Some(((self.name, e),))

	def __repr__(self) -> str:
		return self.name

@dataclass(frozen=True)
class IgnorePattern:
	"""
	A pattern which always succeeds to match, but binds nothing
	"""

	def binds(self, var: str) -> bool:
		"""
		Test to see if this pattern binds a given variable

		For an `IgnorePattern` this is always false
		"""
		return False

	def match(self, e: Expression) -> Option[Sequence[Tuple[str, Expression]]]:
		"""
		Match an expression against this pattern

		>>> IgnorePattern().match(Int(1))
		Some(())
		"""
		return Some(tuple())

	def __repr__(self) -> str:
		return '_'

@dataclass(frozen=True)
class IntPattern:
	value: int

	def binds(self, var: str) -> bool:
		"""
		Test to see if this pattern binds a given variable

		For an `IntPattern` this is always false
		"""
		return False

	def match(self, e: Expression) -> Option[Sequence[Tuple[str, Expression]]]:
		"""
		Match an expression against this pattern

		>>> IntPattern(2).match(Int(1)) is None
		True

		>>> IntPattern(1).match(Int(1))
		Some(())
		"""
		match e:
			case Int(v) if v == self.value:
				return Some(tuple())
		return None

	def __repr__(self) -> str:
		return repr(self.value)

@dataclass(frozen=True)
class SPattern:
	pred: Pattern

	def binds(self, var: str) -> bool:
		"""
		Test to see if this pattern binds a given variable
		"""
		return self.pred.binds(var)

	def match(self, e: Expression) -> Option[Sequence[Tuple[str, Expression]]]:
		"""
		Match an expression against this pattern

		>>> SPattern(NamePattern('n')).match(Int(1))
		Some((('n', 0),))

		>>> SPattern(NamePattern('n')).match(Int(0)) is None
		True

		>>> SPattern(SPattern(NamePattern('n'))).match(Int(4))
		Some((('n', 2),))
		"""
		match e:
			case Int(v) if v > 0:
				return self.pred.match(Int(v - 1))
		return None

	def __repr__(self) -> str:
		return 'S ' + repr(self.pred)

@dataclass(frozen=True)
class Builtin:
	name: str
	f: Callable[[Expression], Option[Expression]]

	def subst(self, expression: Expression, variable: str) -> Expression:
		return self

	def is_value(self) -> bool:
		return True

	def step(self) -> Option[Expression]:
		return None

	def try_apply(self, v: Expression) -> Option[Expression]:
		return self.f(v)

	def __repr__(self) -> str:
		return "'" + repr(self.name)[1:-1] + "'"

	@cur2
	@staticmethod
	def _PLUS_CONST(i: int, e: Expression) -> Option[Expression]:
		match e:
			case Int(v):
				return Some(Int(i + v))
		return None

	@staticmethod
	def _PLUS(e: Expression) -> Option[Expression]:
		match e:
			case Int(v):
				return Some(Builtin(f'+{v}', Builtin._PLUS_CONST(v)))
		return None

	@staticmethod
	def PLUS() -> 'Builtin':
		return Builtin('+', Builtin._PLUS)

	@staticmethod
	def S() -> 'Builtin':
		return Builtin('S', Builtin._PLUS_CONST(1))

@dataclass(frozen=True)
class Function:
	forms: Sequence[Tuple[Pattern, Expression]]

	def subst(self, expression: Expression, variable: str) -> Expression:
		return Function([
			(p, e if p.binds(variable) else e.subst(expression, variable))
			for (p, e) in self.forms
		])

	def is_value(self) -> bool:
		return True

	def step(self) -> Option[Expression]:
		return None

	def try_apply(self, v: Expression) -> Option[Expression]:
		match tuple((bindings.val, body) for (pattern, body) in self.forms for bindings in (pattern.match(v),) if bindings is not None):
			case []:
				return None
			case [(bindings, body), *rest]:
				return Some(subst_all(bindings, body.subst(self, 'recur')))
		raise Exception('Unreachable')

	def __repr__(self) -> str:
		return '{ ' + ', '.join('"' + repr(repr(p))[1:-1] + '" : ' + repr(e) for (p, e) in self.forms) + ' }'

@dataclass
class Application:
	expressions: Sequence[Expression]

	def subst(self, expression: Expression, variable: str) -> Expression:
		return Application([
			e.subst(expression, variable)
			for e in self.expressions
		])

	def is_value(self) -> bool:
		return not len(self.expressions)

	def step(self) -> Option[Expression]:
		match self.expressions:
			case []:
				return None
			case [e]:
				return Some(e)
			case [f, a, *rest]:
				if f.is_value():
					if a.is_value():
						if isinstance(f, Function) or isinstance(f, Builtin):
							return map_opt(
								lambda maybe_f_sub: Application([maybe_f_sub, *rest]),
								f.try_apply(a)
							)
						else:
							return None
					else:
						return map_opt(
							lambda next_a: Application([f, next_a, *rest]),
							a.step()
						)
				else:
					return map_opt(
						lambda next_f: Application([next_f, a, *rest]),
						f.step()
					)
		raise Exception('Unreachable')

	def __repr__(self) -> str:
		return '[ ' + ', '.join(repr(e) for e in self.expressions) + ' ]'

@dataclass
class Int:
	value: int

	def subst(self, expression: Expression, variable: str) -> Expression:
		return self

	def is_value(self) -> bool:
		return True

	def step(self) -> Option[Expression]:
		return None

	def __repr__(self) -> str:
		return str(self.value)

@dataclass
class Variable:
	name: str

	def subst(self, expression: Expression, variable: str) -> Expression:
		if variable == self.name:
			return expression
		else:
			return self

	def is_value(self) -> bool:
		return False

	def step(self) -> Option[Expression]:
		match self.name:
			case '+':
				return Some(Builtin.PLUS())
			case 'S':
				return Some(Builtin.S())
		return None

	def __repr__(self) -> str:
		return '"' + repr(self.name)[1:-1] + '"'

def subst_all(bindings: Sequence[Tuple[str, Expression]], body: Expression) -> Expression:
	match bindings:
		case []:
			return body
		case [(var, replacement), *rest]:
			return subst_all(rest, body.subst(replacement, var))
	raise Exception('Unreachable')