Forem: Gabi

Driving Complete and Easy Bidirectional Typechecking for Higher-Rank Polymorphism in Rust

Gabi — Tue, 15 Aug 2023 01:03:38 +0000

The main goal of this article is making some comments about mb64 implementation of
the Complete and Easy.. paper, but implementing it in pure
rust code with some optimizations, like de bruijin levels and indexes!

Credits

mb64 for this implementation of "Complete and Easy" paper
Sofia for her own implementation of this paper too

Mentions

Practical type inference for arbitrary-rank types

Problems with the paper

The paper uses an ordered context, which is pretty slow to deal with in real world scenarios, and we should get rid of it.

This presupposes that we should use a list of tuples, which is pretty slow to deal with, and we should use a
different data structure, like a hashmap, which is pretty fast to deal with, and we can use the im_rc crate for

Nominal typing, nominal typing makes harder to unify forall expression, and this tutorial presents an idea implementing debruijin indexes and levels

Required Dependencies

The main dependencies are going to be the crates:

thiserror for full-featured error handling in rust
im_rc for immutable data structures with rc
ariadne for presenting language's errors in the stderr or stdout
fxhash for a fast hash implementation

Cool readings

I have made a tutorial about Hindley Milner without any extension, if you don't know anything about type systems, I
highly recommend you to read the following:

My tutorial on Hindley Milner which implements Algorithm W and some another useful things for PL theory in Kotlin.
This response on PL theory stackexchange which explains important stuff type system's notation

Bidirectional Type Systems

Bidirectional type checking is a quite popular technique for implementing type systems with complex features, like
Higher-Rank Polymorphism itself, among other features, like Dependent Type Systems. It's constituted of synthesizing a
type and, checking it against another for comparing two types.

For example, the following expression: 10 : Int, the number 10 synthesizes the type Int, and the type
annotation checks the synthesized type Int against Int specified, and if the wrong type was specified,
like 10 : String, obviously, 10 can't have type String, so the compiler will try to check 10 against String,
and fail at this.

Hindley-Milner

This type checker extends Hindley-Milner implementation, I suggest you read
the Wikipedia's page about hindley milner
implementation.

Our implementation relies on the base of Algorithm J, which is mutable and pretty fast, but it's not pure.

The Hindley-Milner rules we are going to use are the "generalization" and the "unification", to first, create types
without any type annotation, and second and respectively to check if two types are equal, and if they are not, we should
unify them.

Higher-Rank Polymorphism

Higher-Rank Polymorphism is a feature that allows us to write polymorphic functions as arguments, like the following:

let id : ∀ 'a. 'a -> 'a = ... in
let f : (∀ 'a. 'a -> 'a) -> Int = ... in
f id

Pseudo language example which we can pass a polymorphic function as argument

This type of polymorphism is called "Rank-2 Polymorphism", because we are passing in a polymorphic function as argument,
it's different from a naive hindley milner implementation, because we can't pass polymorphic functions as arguments in
hindley milner, because it's not a higher-rank polymorphic type system:

let id : ∀ 'a. 'a -> 'a = ... in
let f : (∀ 'a. 'a -> 'a) -> Int = ... in
f id

The let f part is not valid in a naive hindley milner implementation, because we can't pass polymorphic functions as
arguments

Unification and Subsumption

Unification is the process of checking if two types are equal, and if they are not, we should unify them, for example,
the following expression: 10 : Int, the number 10 synthesizes the type Int, and the type annotation checks the
synthesized type Int against Int specified, and if the wrong type was specified, like 10 : String, obviously, 10
can't have type String, so the compiler will try to check 10 against String, and fail at this.

But "Subsumption" is another term we are going to use, it's the process of checking if a type is a subtype of another,
and the process is called "polymorphic subtyping", for example, the following expressions:

let k : ∀ 'a. ∀ 'b. a -> b -> b = ... in
let f : (Int -> Int -> Int) -> Int = ... in
let g : (∀ 'a. 'a -> 'a) -> Int = ... in
(f k, g k)

Pseudo language example

So, the application (f k) is valid in any ML language, or haskell, etc... Because they are the same type when unified,
but when trying to apply (g k), the compiler will try to unify ∀ 'a. 'a -> 'a with Int -> Int -> Int, and fail at
this, so we need to implement the so-called "subsumption".

The type of ∀ 'a. 'a -> 'a is more "polymorphic" than ∀ 'a. ∀ 'b. a -> b -> b, so the expressions f k and g k
should be valid.

De Bruijin Indexes and Levels

De Bruijin indexes and levels are a way to represent type variables in a type system, for example, the following
expression:

let id : ∀ 'a. a -> a = ... in
id 10

The variable 'a should be replaced by the de bruijin index 0, which is the first variable in the scope

This technique is used to create rigid type variables in some scopes, like if we have a "hole" of a type, and we
increase the de bruijin level of the scope, we won't be able to substitute the variable with another value in our
internal code, for example, the following expression.

The algorithm

The algorithm is based on the paper, but with some modifications, like the de bruijin indexes and levels, and some
optimizations, like the im_rc crate, which is a immutable data structure with rc pointers, and the fxhash crate,
which is a fast hash implementation.

The full code implementation is here.

GADT-like types in Rust

Gabi — Fri, 11 Aug 2023 16:58:42 +0000

GADT-like types in Rust

I think that GADTs are a very powerful feature of Haskell, and I would like to
have something similar in Rust. I think this is the closest thing to GADTs in Rust.

GADTs are the generalized version of algebraic data types (ADTs or enums). In
ADTs, the type of the constructor is the same as the type of the data. In GADTs,
the type of the constructor is more general than the type of the data. This
allows us to express more complex types. In haskell, we can define a GADT like

-- | A GADT for expressions with integers and booleans
-- and addition and equality operations with integers.
data Expr a where
  EInt :: Int -> Expr Int
  EBool :: Bool -> Expr Bool
  EAdd :: Expr Int -> Expr Int -> Expr Int
  EEq :: Expr Int -> Expr Int -> Expr Bool

In Rust, isn't possible to define this kind of types. The thing we can do, is
to define an algebraic data-type:

enum Expr {
  EInt(i32),
  EBool(bool),
  EAdd(Box<Expr>, Box<Expr>),
  EEq(Box<Expr>, Box<Expr>),
}

But we can't specify the type of the constructor. In Haskell, we can use the
kinds to specify the type of the constructor.

We have some kind of implementation already with
the refl crate, but they are
very limited, and verbose. The idea of this post, is showing a simpler and
idiomatic way to implement GADT-like types in rust, just like the refl crate.

The problem

Why do we need GADT? Let's suppose that we want a type that have interior
mutability, but we want to possibly erase the mutability, just like, we have
a Rc<RefCell<..>> and we want to unwrap the RefCell and get a cloned version
of the value in a complex data structure? Just see an example:

enum Expr {
  /// Represents a variable in the source code that can be replaced
  /// by a value if the name is really defined.
  EVar(String, Rc<RefCell<Option<Expr>>>),
  EInt(i32),
  EAdd(Box<Expr>, Box<Expr>),
}

You can take a look in a gist that have more complex data structures
here.

But why do we want to erase the RefCell? Because we want to have Sync, and
Send implementations for the type, and neither RefCell, neither Rc are
Sync or Send. We could have used Arc<Mutex<...>> or Arc<RwLock<...>> but
it would be more verbose, and we would need to use lock() every time we want
to access the value; and if we need a more complex type, just like a closure?
How we would implement it? We can't use Arc<Mutex<...>> because it's a
closure.

The solution

The solution is using traits, and associated types. We can define a trait that
represents the type of the constructor, and then, we can use associated types
to define the type of the data. Let's see an example:

trait State {
  type Variable;
}

We can define a trait that represents the type of EVar constructor. Now, we
can define the type of the data:

enum Expr<S: State> {
  /// Represents a variable in the source code that can be replaced
  /// by a value if the name is really defined.
  EVar(String, S::Variable),
  EInt(i32),
  EAdd(Box<Expr>, Box<Expr>),
}

The S::Variable can be implemented by any type, and we can use it to define
with Rc<RefCell<Option<Expr>>>, and with just Expr if we want to erase the
RefCell and Rc, to be Sync, and Send!

The implementation

We can implement the State trait for any type, and we can define for example,
two states, one state has mutability as previously said, and the other state
has no mutability. Let's see the implementation:

struct MutableState;

impl State for MutableState {
  type Variable = Rc<RefCell<Option<Expr<Self>>>>;
}

struct ImmutableState;

impl State for ImmutableState {
  type Variable = Expr<Self>;
}

And that's it! We can use the Expr<MutableState> and Expr<ImmutableState>
to represent the two different types of expressions. We can also define a
function that converts from Expr<MutableState> to Expr<ImmutableState>. To
be
able to do that, we need to define a trait that converts
from Expr<MutableState>.

Other use cases

There some another cases that traits and GATs can be useful for simulating
GADTs:

Another use case that this kind of implementation is very useful, is when we want to define HOAS (Higher-Order Abstract Syntax) types, and create a dependently-typed system.

HOAS can be very useful when we are building type systems, an example is
Lura's type system, that defines two
states, one state for HOAS, and another state for the Quoted types.

You can have a full look at the
implementation here.
The case of HOAS is very similar to the case of the Expr type, but instead
I need to use closures and Rc.

Conclusion

The conclusion is the following: we can use traits and associated types to
simulate GADTs in Rust. But not fully, but is already very powerful when we have
some complex types that we want to implement, and we want to erase some
information, just like the RefCell and Rc in the Expr type.

Thank you for reading! :)

Writing an Equation Solver

Gabi — Sun, 06 Aug 2023 21:23:46 +0000

Writing an Equation Solver

Writing an Equation Solver is a process that is made of: parsing,
equating/unifying and rewriting.

Equating: it's the step that two terms are equalized and
tries to make equity between them 2, just like an equation
Rewriting: to write equations in real life we need to rewrite
the equation until the variables are discovered, right? For an example:

10 = x + 7
10 - 7 = x + 7 - 7
3 = x
x = 3

So, there was 3 rewrites before the final result. The same process is
applied in a computer, we need to apply some mathematical properties and
rewrite it, and normalise it.

Parsing: we need to get the string of the equation and translate into objects like the following enum:

/// A term in the calculator. This is the lowest level of the calculator.
#[derive(Default, Debug, Clone, Hash, PartialEq, Eq)]
pub enum Expr {
    #[default]
    Error,
    Number(usize),
    Group(Term),
    Variable(Variable),
    BinOp(BinOp),
}

Sections

Mathematical Properties
Parser
- Parser combinators
- Lexer
- Tree
- Parser Implementation
Reducing terms
Symmetry
Distributivity
Associativity
Rewriting
Unifiying/Equating
Final

Mathematical Properties

To write a equation solver, we need to clarify what are the mathematical properties
we are working with. The main ones that we are going to use are the
Associativity, Identity, Commutativity, Symmetry, Distribitivity. We can
define the following rules:

Associativity: (a + b) + c = a + (b + c)
Identity: a = a
Commutativity: a + b = b + a
Symmetry: a + b = c; a = c - b

Ok, right, these mathematical properties are hard to read if we don't have any
code to parallel with it, so let's start writing the parser.

Parser

We need first to tokenize the inputs into a bunch of tokens, which are a kind
of letter with spaces, and some characters ignored:

/// The token type used by the lexer.
#[derive(Debug, PartialEq, Clone)]
pub enum Token<'src> {
    Number(usize),
    Decimal(usize, usize),
    Identifier(&'src str),
    Ctrl(char),
}

The token tooks the lifetime src, because it's referring directly the part
of the source code.

Parser Combinators

We are using technique called parser combinator. And we are using a library chumsky to write parser combinators.

Lexer

The code of the lexer is simply, removing the junk and returning the letters and
numbers as Token. Read the snippet

/// Parse a string into a set of tokens.
///
/// This function is a wrapper around the lexer and parser, and is the main entry point
/// for the calculator.
fn lexer<'src>() -> impl Parser<'src, &'src str, Vec<(Token<'src>, Span)>, LexerError<'src>> {
    let num = text::int(10)
        .then(just('.').then(text::int(10)).or_not())
        .try_map(|(int, decimal), span| {
            // int: &str, decimal: Option<(char, &str)>
            // define the types of the variables, because the chumsky
            // parser tries to infer it
            let int: &str = int;
            let decimal: Option<(char, &str)> = decimal;

            let Ok(int) = int.parse::<usize>() else {
                return Err(Rich::custom(span, "invalid integer"));
            };
            let Some((_, decimal)) = decimal else {
                return Ok(Token::Number(int));
            };

            let Ok(decimal) = decimal.parse::<usize>() else {
                return Err(Rich::custom(span, "invalid decimal"));
            };

            Ok(Token::Decimal(int, decimal))
        })
        .labelled("number");

    // Maps the common mathematical operations like addition, multiplication
    // subtraction, division, and go on..
    let op = one_of("+*-/!^|&<>=")
        .repeated()
        .at_least(1)
        .map_slice(Token::Identifier)
        .labelled("operator");

    // Maps simple incognito variables into identifiers, these are the variables
    // we are trying to discover :)
    let ident = text::ident().map(Token::Identifier).labelled("icognito");

    // The groups that change the precedence.
    let ctrl = one_of("()[]{}").map(Token::Ctrl).labelled("ctrl");

    // Now this finishes the lexer
    num.or(op)
        .or(ctrl)
        .or(ident)
        .map_with_span(|token, span| (token, span))
        .padded()
        .repeated()
        .collect()
}

The full source code is wrote in main.rs.

Tree

We need to define an abstract syntax tree, to translate the mathematical terms
into it:

/// A term in the calculator. This is the lowest level of the calculator.
#[derive(Default, Debug, Clone, Hash, PartialEq, Eq)]
pub enum Expr {
    #[default]
    Error,
    Number(usize),
    Group(Term),
    Variable(Variable),
    BinOp(BinOp),
}

Define the variables and incognitos:

/// A variable. This is a variable that can be assigned to.
#[derive(Default, Debug, Clone)]
pub struct Variable {
    pub name: String,

    /// We use Rc of RefCell here so we can clone and use internal
    /// mutability to change it's value
    pub data: Rc<RefCell<Option<Term>>>,
}

impl Variable {
    /// Gets the data of the variable.
    pub fn data(&self) -> Option<Term> {
        *self.data.borrow().deref()
    }
}

impl PartialEq for Variable {
    fn eq(&self, other: &Self) -> bool {
        self.name == other.name
    }
}

impl Eq for Variable {}

impl Hash for Variable {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.name.hash(state);
    }
}

And finally, the binary operations:

/// A binary operation. This is a binary operation that takes two arguments.
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct BinOp {
    pub op: Op,
    pub lhs: Term,
    pub rhs: Term,
}

The full source code can be found in the main.rs.

Parser implementation

Now we need to write an expression parser, which will take translate tokens into
mathematical terms.

recursive(|expr| {
    // Defines the parser for the value. It is the base of the
    // expression parser.
    let value = select! {
            Token::Number(number) => Expr::Number(number),
            Token::Identifier(identifier) => Expr::Variable(Variable { name: identifier.into(), data: Rc::default() }),
        }
        .map_with_span(|kind, span| (kind, span))
        .map_with_state(|term, _, state: &mut TermArena| state.intern(term))
        .labelled("value");
});

Note thate we used state in the parser, and this is the "arena". The arena will
store the real expressions, and will return an ID to the expression. We will have
to define the Term:

#[derive(Default, Debug, Clone, Copy, Hash, PartialEq, Eq)]
pub struct Term(usize);

And define the "arena":

#[derive(Default)]
pub struct TermArena {
    pub id_to_slot: HashMap<Term, Rc<Spanned<Expr>>>,
    pub slot_to_id: HashMap<Expr, Term>,
}

impl TermArena {
    /// Creates a new term in the arena
    pub fn intern(&mut self, term: Spanned<Expr>) -> Term {
        let id = Term(fxhash::hash(&term.0));
        self.id_to_slot.insert(id, Rc::new(term.clone()));
        self.slot_to_id.insert(term.0, id);
        id
    }

    /// Checks if the term exists in the arena.
    pub fn exists(&self, term: &Expr) -> Option<Term> {
        self.slot_to_id.get(term).cloned()
    }

    /// Gets the term from the arena.
    pub fn get(&self, id: Term) -> Rc<Spanned<Expr>> {
        self.id_to_slot
            .get(&id)
            .cloned()

            // It will return a default expression so the program doesn't crash
            // when the id is missing, it's called sentinel values
            .unwrap_or_else(|| Rc::new((Expr::Error, Range::<usize>::default().into())))
    }
}

Here it's used the library fxhash to fast
hash the values

The arena stuff makes interning of expressions, which is returning the same ID
to same-hash expressions, if the expressions have the same hash, they will return
the same hash, so they will return the same ID. This is useful to improve
performance, and later doing the Commutativvity rule, since 1 + 2 is the
same as 2 + 1, and the hash of it is the same.

So now, we need to write the group terms, like [], {} or even (), so we are
writing the following code in the recursive function.

This is why we need to use recursive, the code will refer to expression

let brackets = expr
    .clone()
    .delimited_by(just(Token::Ctrl('[')), just(Token::Ctrl(']')));

let parenthesis = expr
    .clone()
    .delimited_by(just(Token::Ctrl('(')), just(Token::Ctrl(')')));

let braces = expr
    .clone()
    .delimited_by(just(Token::Ctrl('{')), just(Token::Ctrl('}')));

And we define the "primary" expression, which will catch all values, and
group expressions.

// Defines the parser for the primary expression. It is the
// base of the expression parser.
let primary = value
    .or(braces)
    .or(parenthesis)
    .or(brackets)
    .labelled("primary");

Now we need to define mathematical operations like addition, subtraction,
multiplication and division.

They are splitted in two declarations, so we can have precedence.

let factor = primary
    .foldl_with_state(
        just(Token::Identifier("+"))
            .to(Op::Add)
            .or(just(Token::Identifier("-")).to(Op::Sub))
            .then(expr.clone())
            .repeated(),
        |lhs: Term, (op, rhs), state: &mut TermArena| {
            let (_, fst) = &*state.get(lhs);
            let (_, snd) = &*state.get(rhs);

            let span = SimpleSpan::new(fst.start, snd.end);
            let expr = Expr::BinOp(BinOp { op, lhs, rhs });

            // RULE: Commutativity
            //
            // The commutativity rule states that the order of the operands
            // does not matter. This means that `1 + 2` is the same as `2 + 1`.
            //
            // This rule is implemented by checking if the expression already
            // exists in the state. If it does, then we return the existing
            // expression, otherwise we create a new one.
            match state.exists(&Expr::BinOp(BinOp { op, lhs: rhs, rhs: lhs })) {
                Some(term) if op == Op::Add => term,
                None => state.intern((expr, span)),
                _ => state.intern((expr, span)),
            }
        },
    )
    .labelled("factor");

let term = factor
    .foldl_with_state(
        just(Token::Identifier("*"))
            .to(Op::Mul)
            .or(just(Token::Identifier("/")).to(Op::Div))
            .then(expr.clone())
            .repeated(),
        |lhs: Term, (op, rhs), state: &mut TermArena| {
            let (_, fst) = &*state.get(lhs);
            let (_, snd) = &*state.get(rhs);

            let span = SimpleSpan::new(fst.start, snd.end);
            let expr = Expr::BinOp(BinOp { op, lhs, rhs });

            // RULE: Commutativity
            //
            // The commutativity rule states that the order of the operands
            // does not matter. This means that `1 + 2` is the same as `2 + 1`.
            //
            // This rule is implemented by checking if the expression already
            // exists in the state. If it does, then we return the existing
            // expression, otherwise we create a new one.
            match state.exists(&Expr::BinOp(BinOp { op, lhs: rhs, rhs: lhs })) {
                Some(term) if op == Op::Mul => term,
                None => state.intern((expr, span)),
                _ => state.intern((expr, span)),
            }
        },
    )
    .labelled("term");

The commutativity part starts here, the parser will check if the expression
already exists in the context, and will return the same id, if it exists

Of course, this isn't an automatic process, so we need to reverse the operators
and try to find it as reversed, and if it does exist, we take it from the
context.

Now, we have finished the expression parser, and you can have a look in the
full source code here.
We must write the equation parser, and it's quite simple to do it.

// Defines the parser for the equation. It is the base of the
// parser of equations and inequations.
expr_parser
    .clone()
    .then(
        // Parses an operation
        just(Token::Identifier("="))
            .to(EqKind::Eq)
            .or(just(Token::Identifier("!=")).to(EqKind::Neq)),
    )
    .then(expr_parser.clone())
    .map(|((lhs, op), rhs)| Equation { kind: op, lhs, rhs })
    .map_with_span(|equation, span| (equation, span))
    .labelled("equation")

Its the combination of expr (== | !=) expr.

Reducing terms

We need to reduce the terms to it's normal form, to be compared with another
terms, like: 1 + 2 needs to be reduced to 3 to be compared with 3 properly.

It's the first rewrite rule we need to write! So let's create a function like
rewrite in

impl Term {
    /// Rewrites a term to its normal form.
    pub fn rewrite(self, state: &mut TermArena) -> Term {
        self.normalize(state)
    }
}

And start writing the normalize function.
which is the reduced/or evaluated form.

The link is appointing to the Haskell documentation.

impl Term {
    /// Reduces a term to its normal form.
    pub fn normalize(self, state: &mut TermArena) -> Term {
        let (kind, span) = &*state.get(self);
        let new_kind = match kind {
            Expr::Group(group) => Expr::Group(group.rewrite(state)),
            Expr::BinOp(bin_op) => {
                let lhs = bin_op.lhs.rewrite(state);
                let rhs = bin_op.rhs.rewrite(state);

                // If the term is a number, we try to reduce it evaluating
                // the operation.
                match &state.get(lhs).0 {
                    // If the term is a number
                    //
                    // We assume that the term is a number and we try to
                    // reduce it.
                    Expr::Number(lhs) => match &state.get(rhs).0 {
                        Expr::Number(rhs) => {
                            let number = match bin_op.op {
                                Op::Add => lhs + rhs,
                                Op::Sub => lhs - rhs,
                                Op::Mul => lhs * rhs,
                                Op::Div => lhs / rhs,
                            };

                            Expr::Number(number)
                        }
                        Expr::Group(group) => return group.rewrite(state),
                        _ => return self,
                    },
                    Expr::Group(group) => return group.rewrite(state),
                    _ => return self,
                }
            }
            // If the term is a variable, we try to reduce it.
            Expr::Variable(hole) => match hole.data() {
                Some(value) => return value.rewrite(state),
                None => kind.clone(),
            },
            _ => kind.clone(),
        };
        state.intern((new_kind, *span))
    }
}

It basically evaluates the term to it's normal form

Symmetry

We need to write the symmetry rule, which will be used in a further step called
unifying/equating.

impl BinOp {
    /// RULE: Symmetry
    ///
    /// Examples:
    /// `a + b = c`
    /// `a = c - b`
    pub fn symmetry(&self, value: Term, state: &mut TermArena) -> (Term, Term) {
        let reverse_op = match self.op {
            Op::Add => Op::Sub,
            Op::Sub => Op::Add,
            Op::Mul => Op::Div,
            Op::Div => Op::Mul,
        };

        let span: Span = (state.get(self.lhs).1.start..state.get(self.rhs).1.end).into();

        let lhs = state.intern((
            Expr::BinOp(BinOp {
                op: reverse_op,
                lhs: value,
                rhs: self.rhs,
            }),
            span,
        ));

        (lhs.rewrite(state), self.lhs.rewrite(state))
    }
}

The symmetry is basically, the following steps with the given [BinOp]:

x + 7 = 10
x = 10 - 7

It's a fundamental step for solving equations!

Distributivity

The distributivity property is when we apply an operation to it's left side, like

(x + 2) * 2

Which will be rewrote into:

2 * x + 4

This is fundamental to write some equations, the source code for this step is:

impl Term {
    /// Distributes a term over another term.
    pub fn distribute(self, op: Op, another: Term, state: &mut TermArena) -> Term {
        let (kind, span) = &*state.get(self);
        let new_kind = match kind {
            Expr::Group(_) => return self,
            Expr::BinOp(bin_op) => {
                let lhs = bin_op.lhs.apply_distributive_property(state);
                let rhs = bin_op.rhs.apply_distributive_property(state);

                Expr::BinOp(BinOp {
                    op: bin_op.op,
                    lhs,
                    rhs,
                })
            }
            _ => Expr::BinOp(BinOp {
                op,
                lhs: self,
                rhs: another,
            }),
        };

        state.intern((new_kind, *span))
    }

    /// Applies the distributive property to a term.
    pub fn apply_distributive_property(self, state: &mut TermArena) -> Term {
        let Expr::BinOp(bin_op) = &state.get(self).0 else {
            return self;
        };

        match (&state.get(bin_op.lhs).0, &state.get(bin_op.rhs).0) {
            (Expr::Group(group), _) => group.distribute(bin_op.op, bin_op.rhs, state),
            (_, Expr::Group(group)) => group.distribute(bin_op.op, bin_op.lhs, state),
            (_, _) => self,
        }
    }
}

Associativity

The associativity rules represents an equation that's like: (1 + 2) + 3 = 1 + (2 + 3). It does represents an reorder based in the precedence, to normalize the operation.

We can write a code to associate 3 terms based on it's operator precedence:

/// Applies the associativity rule to a binary operation.
fn associate(lhs: Term, fop: Op, mhs: Term, sop: Op, rhs: Term, state: &mut TermArena) -> BinOp {
    // RULE: Associativity
    //
    // If the term is a binary operation, we try to reduce it
    // to its weak head normal form using the precedence of
    // the operators.
    //
    // This step is called precedence climbing.
    //
    // Evaluate the operation if the precedence of the
    // operator is higher than the precedence of the
    // operator of the right hand side.
    let lhs = lhs.apply_associativity(state).rewrite(state);
    let mhs = mhs.apply_associativity(state).rewrite(state);
    let rhs = rhs.apply_associativity(state).rewrite(state);

    // If the precedence of the operator of the left hand side
    // is higher than the precedence of the operator of the
    // right hand side, we change the order.
    if op_power(sop) >= op_power(fop) {
        BinOp {
            op: fop,
            lhs,
            rhs: state.intern((
                Expr::BinOp(BinOp {
                    op: sop,
                    lhs: mhs,
                    rhs,
                }),
                (0..0).into(),
            )),
        }
    } else {
        BinOp {
            op: fop,
            lhs: state.intern((
                Expr::BinOp(BinOp {
                    op: sop,
                    lhs,
                    rhs: mhs,
                }),
                (0..0).into(),
            )),
            rhs,
        }
    }
}

And write a wrapper in [Term] to call it with [BinOp] operations:

impl Term {
    /// Applies the associativity rule to a term.
    pub fn apply_associativity(self, state: &mut TermArena) -> Term {
        let (kind, span) = &*state.get(self);

        // If the term is not a binary operation, we return it.
        let Expr::BinOp(mut bin_op) = kind.clone() else {
            return self;
        };

        // Apply associativy to the leftmost side of the expression.
        //
        // This is done by recursively applying the associativity
        if let Expr::BinOp(lhs_bin) = &state.get(bin_op.lhs).0 {
            bin_op = associate(
                lhs_bin.lhs,
                lhs_bin.op,
                lhs_bin.rhs,
                bin_op.op,
                bin_op.rhs,
                state,
            );
        }

        // Apply associativy to the rightmost side of the expression.
        //
        // This is done by recursively applying the associativity
        if let Expr::BinOp(rhs_bin) = &state.get(bin_op.rhs).0 {
            bin_op = associate(
                bin_op.lhs,
                bin_op.op,
                rhs_bin.lhs,
                rhs_bin.op,
                rhs_bin.rhs,
                state,
            );
        }

        // Reintern the term.
        state.intern((Expr::BinOp(bin_op), *span))
    }
}

Rewriting

We need to change the rewrite function to compute all rewrite rules we have
made:

/// Rewrites a term to its normal form.
pub fn rewrite(self, state: &mut TermArena) -> Term {
    self.apply_distributive_property(state)
        .apply_associativity(state)
        .normalize(state)
}

Unifiying/Equating

Now we need to write the logical part, which will compare the operations. We need
to first start a pattern matching:

impl Term {
    /// Unifies two terms. It's the main point of the equation.
    ///
    /// The logic relies here.
    pub fn unify(self, another: Term, state: &mut TermArena) -> Result<(), TypeError> {
        match (&state.get(self).0, &state.get(another).0) {
            // Errors are sentinel values, so they are threated like holes
            // and they are ignored, anything unifies with them.
            (Expr::Error, _) => {}
            (_, Expr::Error) => {}

            (a, b) => {
                return Err(TypeError::NotUnifiable(a.clone(), b.clone()));
            }
            // ...
        }

        Ok(())
    }
}

The terms successfully unified will fallback ino the Ok(()) expressions,
and if it's not unified, it will return an error just like in the TypeError::NotUnifiable
part

And wrap the group terms, unifying its values:

// Reduce groups to the normal form
(_, Expr::Group(another)) => {
    let term = self.rewrite(state);
    let another = another.rewrite(state);

    term.unify(another, state)?;
}
(Expr::Group(term), _) => {
    let term = term.rewrite(state);
    let another = another.rewrite(state);

    term.unify(another, state)?;
}

And unify numbers, if they are the same number, it will unify properly

// If they are the same, they unify.
(Expr::Number(a), Expr::Number(b)) if a == b => {}

Now, the variable stuff, which is the same as the number part, if the
name is the same, it will unify.


// If they are variables, they unify if they are the same.
//
// This check isn't inehenterly necessary, but it's a good catcher
// to avoid panics, because if the variables are the same, they will
// try to borrow the same data, and it will panic with ref cells.
(Expr::Variable(variable_a), Expr::Variable(variable_b))
    if variable_a.name == variable_b.name => {}

Now we need to unify attributions, like x = 5, or something like this, this
will basically give the variable, a meaning.

// Unifies the variable with the term, if the variable is not bound. If
// it's bound, it will try to unify, if it's not unifiable, it will
// return an error.
(_, Expr::Variable(variable)) => {
    match variable.data() {
        // If the variable is already bound, we unify the bound
        Some(bound) => {
            self.unify(bound, state)?;
        }
        // Empty hole
        None => {
            variable.data.replace(Some(self));
        }
    }
}
(Expr::Variable(variable), _) => {
    match variable.data() {
        // If the variable is already bound, we unify the bound
        Some(bound) => {
            bound.unify(another, state)?;
        }
        // Empty hole
        None => {
            variable.data.replace(Some(another));
        }
    }
}

And now, unify the operations, like 1 + 1 = 1 + 1, we need to reduce the
terms, like: 1 + 1 is equivalent to 2, so we will compare 2 with 2,
and it will successfully unify:

// Unifies the bin ops if they are the same, and unifies the
// operands.
(Expr::BinOp(bin_op_a), Expr::BinOp(bin_op_b)) => {
    if bin_op_a.op == bin_op_b.op {
        let lhs_a = bin_op_a.lhs.rewrite(state);
        let rhs_a = bin_op_a.rhs.rewrite(state);

        let lhs_b = bin_op_b.lhs.rewrite(state);
        let rhs_b = bin_op_b.rhs.rewrite(state);

        lhs_a.unify(lhs_b, state)?;
        rhs_a.unify(rhs_b, state)?;
    } else {
        return Err(TypeError::IncompatibleOp(bin_op_a.op, bin_op_b.op));
    }
}

Now, the last part, we need to unify the symmetry, if we have an example like
the comments, it will use the symmetry function.

// If the term is a number, we try the following steps given the example:
//   9 = x + 6
//
// 1. We get the reverse of `+`, which is `-`.
// 2. We subtract `6` from both sides, to equate it,
//    and since the `9` is a constant, we can reduce
//    it to the normal form.
// 3. Got the equation solved, we can unify the `x` with
//    the result of the subtraction.
//
//    3 = x and then x = 3
(Expr::Number(_), Expr::BinOp(bin_op)) => {
    let (term, another) = bin_op.symmetry(self, state);

    term.unify(another, state)?;
}
(Expr::BinOp(bin_op), Expr::Number(_)) => {
    let (term, another) = bin_op.symmetry(another, state);

    term.unify(another, state)?;
}

Final

Thanks for your read :) Have a nice day!