A walker / visitor for pliron

[urso]

In order to give pliron a test run I'm thinking to start with a very simple language. I chose the kaleidoscope language from the LLVM tutorials. One starts very simple with expressions and functions calls. And later eventually adds some control constructs. This makes it a great way to test things out while having a very small set Ops to deal with.

The idea is parse it into a "high level" dialect modeling the language, then lower it to a cranelift based dialect, convert the IR to cranelift and finally execute the generated code.

This also gives us a means to discuss enhancements to pliron on live (but hopefully simple) code.

For the beginning:

Issue: how to use function arguments with the basic blocks?

Do we have an API or plans for extending the API of FuncOp to register arguments that can be used as Operands in the function body? See #24

Visitor?

We do not have any functionality to walk or rewrite operations, regions, and blocks. There is an draft PR #4, but maybe we should start with a walker API first? Is this something that would fit this project, or shall pliron rather focus on the storage + print/parse part for now while we start pushing more complex functionality like walking/rewriting/lowering/pass-manager into external repos and eventually incorporate some functionality that works well back into pliron?

#4 has a visitor that seems to resemble the MLIR walker functions.

Support to visit operations recursively is an important building block for more complex operations/transformations. Especially with the potentially arbitrary nesting of IR Ops and regions. Maybe we should start with that pattern and see how it goes.

In contrast to MLIR and #4, I have a more typical visitor pattern based on traits in mind. Maybe even separating the 'Walk' from the 'Visit' part.

[vaivaswatha]

I chose the kaleidoscope language

This is definitely a good choice!. Thank you.

then lower it to a cranelift based dialect, convert the IR to cranelift and finally execute the generated code.

I had LLVM in mind (only because I'm more familiar with it), but cranelift is a good choice too. I think at some point we'll have them both. No reason to restrict.

This also gives us a means to discuss enhancements to pliron on live (but hopefully simple) code.

💯

maybe we should start with a walker API first? Is this something that would fit this project,

Yes, very much. I've been pondering over a rewrite infrastructure for some time, but that should begin with a walker API first, so yes.
Regarding rewrites, at least for 1->many rewrites, i.e., replacing one operation with one or more other operations, we could do it with the Interface infrastructure we currently have. If an Op wants to be part of a particular rewrite (say Dialect1ToDialect2), then it implements the Dialect1ToDialect2 interface, which has a rewrite fn that looks like (Op) -> Vec<Op>. so the rewriter (walker?), whenever it finds a particular Op would just check if it implements that particular interface, and if it does, replace the Op with whatever it returns upon calling the rewrite fn.

I have a more typical visitor pattern based on traits in mind. Maybe even separating the 'Walk' from the 'Visit' part.

Can you elaborate a little more? I'd like to hear this.

[urso]

I opened a PR for the function arguments: #26

The change kinda makes me wonder if we want to have a more type-safe API. I think it is not the first time I bring this up.
Right now we pass Ptr<Type> and Ptr<AttribObj> around, which means that we must reinterpret the actual "storage" type every now and then. As a result of this a program can fail at runtime, because we bypass the Rust type-system.
It also makes the API more difficult to use as it is less discoverable which IR types interact with which Op types.

One common pattern we find in MLIR is the separation of the 'storage' from the 'interpretation'. We have it in the Op and Operation types. The Operation provides a standardized storage type for IR ops and the Op and the interfaces give us a type-safe interpretation of the underlying storage value.

In contrast to pliron we find a similar pattern in MLIR for types and attributes. AttributeStorage and TypeStorage types. The storage types do store the actual fields of the structs. The concrete types like IntegerType or FunctionType do not really have any fields, but access the fields via their corresponding storage classes.

Let's look at an example for IntegerType:

struct TypePtr<T: TypeStorage> = PtrOf<TypeObj, T>;

struct PtrOf<Kind, T>(Ptr<Kind>, PhantomMarker<T>);

struct IntegerTypeStorage {
  width: u64,
  signedness: Signedness,
}

struct IntegerType(TypePtr<IntegerTypeStorage>);  // alternatively `struct IntegerType = TypePtr<IntegerTypeStorage>`.

The IntegerTypeStorage can be reused allowing users to give it a different interpretations. This actually gives us a family of concrete types that developers can require/expect in their own Ops/Attributes.

type Int8Type(TypePtr<IntegerTypeStorage>);

type Int16Type(TypePtr<IntegerTypeStorage>);

...

type Unit8Type(...);

type Uint16Type(...);

The split into Type and TypeStorage also means that we would need to rethink some of the related traits:

pub trait Storage {
  fn eq(&self, &dyn Storage) -> bool;

  // hash?
}

pub trait TypeStorage: Storage {
  fn hash(&self) -> TypeHashValue;

  ...
}

pub trait Type {
    type Storage: TypeStorage;

    fn get_type_id(&self) -> TypeId { Self::get_type_id_static() }

    fn get_type_id_static() -> TypeId
    where
        Self: Sized;

    fn get_ptr(&self) -> Ptr<TypeObj> {
      self.get_type_ptr().erase_type()
    }

    fn get_type_ptr(&self) -> TypePtr<Self::Storage> { ... }

    fn from_storage(tptr: TypePtr<Self::Storage>) -> Self;

    fn from_ptr(ptr: Ptr<TypeObj>) -> Option<Self> {
      let typed = ... ?;
      Some(from_storage(typed));
    }
}

Now the constructor of FuncOp would become become:

impl FuncOp {
    pub fn new_unlinked(ctx: &mut Context, name: &str, ty: FuncType) -> FuncOp {
      ....
    }
}

Other Ops like ConstantOp might still allow any kind of Type or Attribute. For this to work we have the get_ptr method that will gives the type erased value.

One can also accept a type family via TypePtr<IntegerTypeStorage>.

[vaivaswatha]

The change kinda makes me wonder if we want to have a more type-safe API

I agree with your concern. Let me think about this a bit and get back to you.

[urso]

Can you elaborate a little more? I'd like to hear this.

Did some experiments and for now settled on the following approach:

The core traits are Acceptor and Strategy. I have reserved Visitor for later. Plus I also introduced a AcceptorMut and StrategyMut trait (also these ones have been a little tricky). As loads of functionality is more or less copy'n paste I will try to focus on a few traits/types only for sake of the discussion.

On top of some core traits I have implemented some more implementations and other helpers to make some common patterns more convenient to use. Kind of like a Lego-set of walking patterns 😅 .

IR walking - core

The Acceptor trait is more or less the actual core trait and users can modify visiting strategy at will by implementing that trait. The acceptor trais also comes with a default implementation of all its methods, so users can selectively change only a subset of the functionality while "walking" for other IR entities use the default implementation.

The Acceptor trait looks as follows:

pub trait Acceptor<B> {
    fn accept_op(&mut self, ctx: &Context, op: Ptr<Operation>) -> WalkResult<B> {
        Forward::walk_regions(self, ctx, op)
    }

    fn accept_region(&mut self, ctx: &Context, region: Ptr<Region>) -> WalkResult<B> {
        Forward::walk_blocks(self, ctx, region)
    }

    fn accept_block(&mut self, ctx: &Context, block: Ptr<BasicBlock>) -> WalkResult<B> {
        Forward::walk_ops(self, ctx, block)
    }
}

The Acceptor itself accepts an IR entity (hence the name), but does not implement any strategy for iterating over child-nodes itself. Instead we delegate the walking part to a Strategy. Custom acceptors are free to reuse existing strategies, or just plainly ignore them.

The Strategy trait looks as follows:

pub trait Strategy {
    fn walk_regions<A, B>(a: &mut A, ctx: &Context, op: Ptr<Operation>) -> WalkResult<B>
    where
        A: Acceptor<B> + ?Sized;

    fn walk_blocks<A, B>(a: &mut A, ctx: &Context, region: Ptr<Region>) -> WalkResult<B>
    where
        A: Acceptor<B> + ?Sized;

    fn walk_ops<A, B>(a: &mut A, ctx: &Context, block: Ptr<BasicBlock>) -> WalkResult<B>
    where
        A: Acceptor<B> + ?Sized;
}

We have 2 implementations: Forward and Reverse (sometimes you might want to walk backwards). The default implementation of Acceptor uses the Forward strategy.

You may have noticed that we return a WalkResult, which uses just a type alias over std::ops::ControlFlow:

pub type WalkResult<B> = ControlFlow<B, WalkCont>;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum WalkCont {
    Advance, // Advance the walk
    Skip,    // Advance to the next sibling
}

As ControlFlow implements the Try trait (which custom types can not implement in stable) we get the ? operator for free.

The ControlFlow::Break value allows users to return data from their walking operation to the caller. This could just be (), some IR object of interest, or an error.

Pre-/Post Order traversal

Having to implement Acceptor all the time might not be that convenient. Instead we would provide a PreOrderAcceptor and PostOrderAcceptor out of the box. Each of them accepts a Visitor. Once again the Visitor has a default implementation for each type, such that users can concentrate on the IR kind they are interested in:

pub trait Visitor<B> {
    fn visit_op(&mut self, _ctx: &Context, _op: Ptr<Operation>) -> WalkResult<B> {
        WalkResult::Continue(WalkCont::Advance)
    }

    fn visit_region(&mut self, _ctx: &Context, _region: Ptr<Region>) -> WalkResult<B> {
        WalkResult::Continue(WalkCont::Advance)
    }

    fn visit_block(&mut self, _ctx: &Context, _block: Ptr<BasicBlock>) -> WalkResult<B> {
        WalkResult::Continue(WalkCont::Advance)
    }
}

The PreOrderAcceptor uses the Visitor and the configured Strategy:

pub struct PreOrderAcceptor<B, S, V>(V, PhantomData<(S, B)>);

impl<B, S: Strategy, V: Visitor<B>> Acceptor<B> for PreOrderAcceptor<B, S, V> {
    fn accept_op(&mut self, ctx: &Context, op: Ptr<Operation>) -> WalkResult<B> {
        let c = self.0.visit_op(ctx, op)?;
        if c == WalkCont::Advance {
            S::walk_regions(self, ctx, op)?;
        }
        WalkResult::Continue(WalkCont::Advance)
    }

    fn accept_region(&mut self, ctx: &Context, region: Ptr<Region>) -> WalkResult<B> {
        let c = self.0.visit_region(ctx, region)?;
        if c == WalkCont::Advance {
            S::walk_blocks(self, ctx, region)?;
        }
        WalkResult::Continue(WalkCont::Advance)
    }

    fn accept_block(&mut self, ctx: &Context, block: Ptr<BasicBlock>) -> WalkResult<B> {
        let c = self.0.visit_block(ctx, block)?;
        if c == WalkCont::Advance {
            S::walk_ops(self, ctx, block)?;
        }
        WalkResult::Continue(WalkCont::Advance)
    }
}

Enter/Exit? Stack based analyzers/processing

Sometimes one wants to be notified when entering and exiting a particular node type. An acceptor supporting said pattern might look like this:

struct EnterExit<B, S, V>(V, PhantomData<(S, B)>);

impl<B, S: Strategy, V: SurroundingVisitor<B>> Acceptor<B> for EnterExit<B, S, V> {
    fn accept_op(&mut self, ctx: &Context, op: Ptr<Operation>) -> WalkResult<B> {
        let c = self.0.enter_op(ctx, op)?;
        if c == WalkCont::Advance {
            S::walk_regions(self, ctx, op)?;
        }
        self.0.exit_op(ctx, op)
    }

    fn accept_region(&mut self, ctx: &Context, region: Ptr<Region>) -> WalkResult<B> {
        let c = self.0.enter_region(ctx, region)?;
        if c == WalkCont::Advance {
            S::walk_blocks(self, ctx, region)?;
        }
        self.0.exit_region(ctx, region)
    }

    fn accept_block(&mut self, ctx: &Context, block: Ptr<BasicBlock>) -> WalkResult<B> {
        let c = self.0.enter_block(ctx, block)?;
        if c == WalkCont::Advance {
            S::walk_ops(self, ctx, block)?;
        }
        self.0.exit_block(ctx, block)
    }
}

Walking with callback functions?

On top of that functionality one can provide callback based walker functions:

pub fn walk_op<B, O: Order, S: Strategy>(
    root_op: Ptr<Operation>,
    ctx: &Context,
    callback: &mut impl FnMut(&Context, Ptr<Operation>) -> WalkResult<B>,
) -> WalkResult<B> {
    let cb = Callback(callback, PhantomData);
    let mut acceptor = O::build::<B, S, _>(cb);
    acceptor.accept_op(ctx, root_op)
}


pub fn walk_regions<B, O: Order, S: Strategy>(
    root_op: Ptr<Operation>,
    ctx: &Context,
    callback: &mut impl FnMut(&Context, Ptr<Region>) -> WalkResult<B>,
) -> WalkResult<B> {
    let cb = Callback(callback, PhantomData);
    let mut acceptor = O::build::<B, S, _>(cb);
    acceptor.accept_op(ctx, root_op)
}

pub fn walk_blocks<B, O: Order, S: Strategy>(
    root_op: Ptr<Operation>,
    ctx: &Context,
    callback: &mut impl FnMut(&Context, Ptr<BasicBlock>) -> WalkResult<B>,
) -> WalkResult<B> {
    let cb = Callback(callback, PhantomData);
    let mut acceptor = O::build::<B, S, _>(cb);
    acceptor.accept_op(ctx, root_op)
}

The Callback wrapper captures the function and provides a visitor implementation as follows:

struct Callback<B, IR, F: FnMut(&Context, Ptr<IR>) -> WalkResult<B>>(F, PhantomData<IR>);

impl<B, F: FnMut(&Context, Ptr<Operation>) -> WalkResult<B>> Visitor<B>
    for Callback<B, Operation, F>
{
    fn visit_op(&mut self, ctx: &Context, op: Ptr<Operation>) -> WalkResult<B> {
        (self.0)(ctx, op)
    }
}

impl<B, F: FnMut(&Context, Ptr<Region>) -> WalkResult<B>> Visitor<B> for Callback<B, Region, F> {
    fn visit_region(&mut self, ctx: &Context, region: Ptr<Region>) -> WalkResult<B> {
        (self.0)(ctx, region)
    }
}

impl<B, F: FnMut(&Context, Ptr<BasicBlock>) -> WalkResult<B>> Visitor<B>
    for Callback<B, BasicBlock, F>
{
    fn visit_block(&mut self, ctx: &Context, block: Ptr<BasicBlock>) -> WalkResult<B> {
        (self.0)(ctx, block)
    }
}

The Order trait acts as an Acceptor-builder.

Example usage of walk_ops:

walk_ops::<_, PreOrder, Forward>(my_op, ctx, |ctx, op| {
  ...
});

&mut Context

Without &mut Context common strategy to rewrite the IR input would be to collect all ops into a set of to-be-rewritten Ops and then iterate over the set in order to update the nodes. In case of lowering we would basically create a Vec<Ptr> of all existing Ops in the input program.

For this reason I have basically copied all interfaces and provide XMut instance and x_mut methods (thank you Copilot).

Unfortunately the borrow checker doesn't like that the Iterator has a borrow on ctx while we pass the &mut Context along. For this reason I have had to introduce a temporary allocation to capture all Ptr<X> into a Vec. For example:

    fn walk_regions_mut<A, B>(a: &mut A, ctx: &mut Context, op: Ptr<Operation>) -> WalkResult<B>
    where
        A: AcceptorMut<B> + ?Sized,
    {
        let iter = { // use a block to capture `ctx` to force a reborrow
            let op = op.deref(ctx);
            op.regions().iter().cloned().collect::<Vec<_>>()
        };
        for ptr in iter {
            a.accept_region_mut(ctx, ptr)?;
        }
        WalkResult::Continue(WalkCont::Advance)
    }

Having many small allocations and free does not sound that nice either. In order to get rid of the allocations one would need Context to support some form of interior mutability I guess. In that case we might not even need the Mut variants.

Follow ups

These are still rather low-level building blocks. Users get to operate on Ptr<Operation>, Ptr<Region>, or Ptr<BasicBlock>. As a next step one might look into some filtering capability or support to declare a callback function to operator on a specific type that implements Op or some other op interfaces for convenience.

You can find the complete code (still untested) here.

The impl Iterator return for regions() didn't work well with reverse iterations and the compiler did complain when I tried to turn this into a DoubleEndedIterator. In order to be more flexible I just changed regions to:

 pub fn regions(&self) -> &[Ptr<Region>] {
        &self.regions
    }

Alternative strategy:

In visitors.rs I starting using a different strategy. Here we start from bottom up and have a `Walker that does the walking and call a visitor that accepts an enter and exit callback. Higher layer reuse that kind of walker as Core and for example ignore the "post-visit" callback. But in the end I did prefer the solution where accept + walking strategy are separated.

[vaivaswatha]

I've started looking at both your proposals. I haven't gone into the details yet, but visitor.rs seems simpler to me. I'm still trying to work my way around the many traits in visitor2.rs. The names especially make it harder to grasp what the functionality is. For example,

pub enum WalkCont {
    Advance, // Advance the walk
    Skip,    // Advance to the next sibling
}

They're both advancing the walk, but one is called Advance while the other is called Skip. I have the same confusions b/w acceptor, visitor etc.

But it's too early to bikeshed on names !. Let me spend a little more time.

[urso]

Maybe ignore visitor2.rs. That one was rather an intermediate step to play with some traits. visitor3.rs is more or less the final result.

I haven't gone into the details yet, but visitor.rs seems simpler to me.

Interesting. Personally I find visitor3 easier to understand because it has a smaller core trait and everything else is builtin around/on top of that core trait. In visitor.rs the core trait is Walker which also requires WalkerVisitor. IMO these traits have a quite "large" interface and also requires implementors to call a pre and post method on some kind of delegate visitor. That is Walker implements everything and then we build simpler interfaces on top of it.

In visitor3.rs the core trais is much simpler. It is the Acceptor trait which makes no real assumptions on how the walk is to be implemented or if a delegate should be used. That's it. Everything else in visitor3.rs is mostly for convenience. I think smaller interfaces allow for better composability.

They're both advancing the walk, but one is called Advance while the other is called Skip.

Yeah, this concept is in each proposal. It is actually the same idea as in #4. Here we have:

/// A utility result that is used to signal how to proceed with an ongoing walk:
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum WalkResult {
    /// the walk will be interrupted and no more operations, regions
    /// or blocks will be visited.
    Interrupt,
    /// the walk will continue.
    Advance,
    /// the walk of the current operation, region or block and their
    /// nested elements that haven't been visited already will be skipped and will
    /// continue with the next operation, region or block.
    Skip,
}

The drawbacks of WalkResult are:

• does not implement the Try trait. So we have to introduce match statements which I found quite noise.
• Can not return a value or error on Interrupt.

The std::ops::ControlFlow is defined as:

pub enum ControlFlow<B, C = [()](https://doc.rust-lang.org/std/primitive.unit.html)> {
    Continue(C),
    Break(B),
}

The ControlFlow::Break behaves similar to WalkResult::Interrupt, but we can return a value and ControlFlow implements the Try trait so we can use ?.

I still think the concepts Advance and Skip from WalkResult are important to have, which is why I introduced WalkCont for the cases where we do not want to abort the walk. I just copied the names. This gives us:

WalkResult::Accept == ControlFlow::Continue(WalkCont::Advance)
WalkResult::Skip == ControlFlow::Continue(WalkCont::Skip)
WalkResult::Interrupt == ControlFlow::Break(())

I would have preferred if we could implement the Try trait on a custom WalkResult<T> as we have it in #4. But unfortunately this is not possible yet in stable (rust-lang/rust#84277). Yet I find the convenience of the ? operator intriguing enough to use ControlFlow as a compromise.

I have the same confusions b/w acceptor, visitor etc.

The traits are "inspired" by a mix of the pattern used in rustc (method for each object type + default implementations + separate reusable walk functions), PR #4 (callback as delegate + WalkResult for control flow).

The Acceptor trait similar to the rustc_ast::visit::Visitor trait. In rustc we have a number of walk_x companion functions that act as default implementations for the Visitor trait and allow users to reuse their implementation in custom traits at will. Instead of a set of functions I introduced the Strategy trait with 2 implementations (Forward and Reverse), which implement the actual walk. They are reusable is in rustc_ast, but they give the user some more flexbility in on how they want the walk to be implemented. Plus being generics we can allow more generic Acceptor implementations to also accept a walking strategy as a parameter.

But in the end the Acceptor trait is the only trait you would need, everything else is for convenience.

In my mind when working with an AST like structure one is interested in 3 simple use cases most the time:

• pre-order-traversal
• post-order-traversal
• enter-exit callbacks (e.g. in case the internal visitor state requires to manage a stack).

For these 3 use-cases we provide default implementations. The PreOrderAcceptor, PostOrderAcceptor and EnterExit (for lack of a better name).

As inspiration from #4 I made these default implementations accept a "delegate" (basically a callback). Because In my mind these patterns are so common that I named the trait for the delegate Visitor. This is also why the actual "visitor" is named Acceptor and we have accept_x methods. The name is chosen to some extend from the common practice in Java to have an accept(Visitor) methods on the concrete elements.

Naming is hard. Would it help if we rename Acceptor to Visitor and the current Visitor to Delegate (or Callback)?

One more remark regarding WalkCont. I don't think we need that one always. The PostOrder based visitor has no need for WalkCont. In that case it would make sense to just require ControlFlow<B> as the return type. Only the PreOrder and EnterExit visitors need the WalkCont as a return value from their delegates.

More simplification to visitor3.rs:

The visitor3.rs also implmements the Mut based traits that are missing from visitor.rs, so I don't think it is that much more code (actually might be less). I'm also thinking to remove the EnterExit acceptor as it is easy enough to just implement some stack like visitor like so:

struct MyStackVisitor {
  stack Vec<State>
}

impl Acceptor<()> for MyStackVisitor {
  fn accept_op(&mut self, ctx: &Context, op: Ptr<Operation>) -> WalkResult<()> {
    // on enter code
    self.stack.push(....)

    Forward::walk_regions(self, ctx, op)?;

    // on exit code
    self.stack.pop()
  }
}

Removing EnterExit would get rid of 2 types and 100 lines of code.

Similarly one might even question the PostOrderAcceptor and PreOrderAcceptor. For example:

struct MyPreOrderVisitor {
}

impl Acceptor<()> for MyStackVisitor {
  fn accept_op(&mut self, ctx: &Context, op: Ptr<Operation>) -> WalkResult<()> {
    ...
    Forward::walk_regions(self, ctx, op)
  }
}

For convenience it would still be nice to have the callback based walk_ops functions, but these maybe could be build without those other implementations.

I think I will give that idea a try the next days. Maybe can simplify the implementation and make things more clear.

[vaivaswatha]

Would it help if we rename Acceptor to Visitor and the current Visitor to Delegate (or Callback)?

Actually it would. But before that, let me take some time to ponder over this.

[urso]

fyi; I added a visitor4.rs that removes some types from visitor3.rs and renames Acceptor to Visitor: https://github.com/urso/pliron/blob/visitor/src/visitor4.rs

This one also adds some tests.

Right now the order and direction are configured via generics. For the Visitor trait this is ok, but for the walk_x functions as used in the tests one might prefer to configure the walk via enums.

[vaivaswatha]

I'm looking at the latest, visitor4.rs. It's much more comprehensible now, thank you very much. Before I get in deeper, I had a fundamental question.

why do we need a separation b/w Visitor and Strategy traits? From what I can tell, we don't need Visitor at all, but just the Strategy.

[urso]

Why do you think we need Strategy only?

The Strategy implements the walking part, while visitor implements the callback part. The visitor can choose on how to process the IR entity or if it wants to continue walking. In OOP this is normally done by having an accept method on the concrete class that can be visited. [This is briefly discussed](https://rust-unofficial.github.io/patterns/patterns/behavioural/visitor.html#discussion in the unofficial Rust patterns book. I only decided to make the 'walk' part configurable by moving the 'walk' functions into a trait.

[vaivaswatha]

After spending some more time on visitor4.rs.

I still find this code quite hard to follow through, mostly the various instantiations. For example, why does walk_blocks, a function that should walk over blocks (of a region?) take a root_op: Ptr<Operation> as its argument? On the other hand, right above there, walk_op does what? The singular name indicates that it walks over one op? or should it have been walk_ops, in which case it must've taken a Ptr<BasicBlock> as its argument, walking over the ops in a block.

Strategy and Order being template parameters everywhere makes it quite hard to grasp the overall flow.

[urso]

I still find this code quite hard to follow through, mostly the various instantiations. For example, why does walk_blocks, a function that should walk over blocks (of a region?) take a root_op: Ptr as its argument? On the other hand, right above there, walk_op does what? The singular name indicates that it walks over one op? or should it have been walk_ops, in which case it must've taken a Ptr as its argument, walking over the ops in a block.

The names are more or less and artifact because I initially started with references to concrete types like &Operation and so on. When doing the VisitorMut I had to switch to the Ptr types to eventually satisfy the borrow checker. I didn't clean up the names and for a proper PR there should be some more cleanup and testing.

I renamed the walk_x methods in the Strategy trait to match the IR entity. Maybe it is more clear now, as it is more similar

More important to me right now is the potential direction/approach to the API like traits/types/utility functions we want to expose.
For example:

Strategy and Order being template parameters everywhere makes it quite hard to grasp the overall flow.

Right.

The code consists of 2 functionalities. The actual visitor core (Visitor and Strategy traits) and the walk_op/walk_region/walk_block utility functions allowing users to work with a callback instead of having to implement the visitor trait for each use case.

One does not really need the second one, but I like it for convenience. The Order trait is only part of the later one. I could also have chosen an Enum for Order and instead of exposing Strategy require an enum for the "direction". Then choose the correct Strategy in the implementation based on the chosen enum values. Like:

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum VisitOrder {
    PreOrder,
    PostOrder,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Direction {
    Forward,
    Reverse,
}

pub fn walk_op_with_enums<B>(
    order: VisitOrder,
    direction: Direction,
    root: Ptr<Operation>,
    ctx: &Context,
    callback: impl FnMut(&Context, Ptr<Operation>) -> WalkResult<B>,
) -> WalkResult<B> {
    match (order, direction) {
        (VisitOrder::PreOrder, Direction::Forward) => {
            walk_op::<B, PreOrder, Forward, _>(root, ctx, callback)
        }
        (VisitOrder::PreOrder, Direction::Reverse) => {
            walk_op::<B, PreOrder, Reverse, _>(root, ctx, callback)
        }
        (VisitOrder::PostOrder, Direction::Forward) => {
            walk_op::<B, PostOrder, Forward, _>(root, ctx, callback)
        }
        (VisitOrder::PostOrder, Direction::Reverse) => {
            walk_op::<B, PostOrder, Reverse, _>(root, ctx, callback)
        }
    }
}

I actually added this implementation to visitor4.rs. I agree here, the enum based implementation looks much nicer in comparison to all the type parameters one would have to fill out otherwise.

Now we do not need to export the Callback visitor. When implementing an actual visitor you would normally also set the Strategy when implementing the trait. The Callback visitor is special as it generalizes over the direction of the walk. So in downstream code we would see custom visitors like so:

struct MyOpAnalyzer{...};

// Walk tree forward in pre order direction. We return an error value when doing a Break.
impl Visitor<Forward, Error> for MyOpAnalyzer {
  fn visit_op(&mut self, ctx: &Context, op: Ptr<Operation>) -> WalkResult<Error> {
    ...
  }
}

Because the other visit method use the configured strategy we do not have to implement them if we want to choose another strategy.

If we change Visitor to this

pub trait Visitor<B=(), S: Strategy = Forward> {
    fn visit_op(&mut self, ctx: &Context, op: Ptr<Operation>) -> WalkResult<B> {
        S::walk_op(self, ctx, op)
    }

    fn visit_region(&mut self, ctx: &Context, region: Ptr<Region>) -> WalkResult<B> {
        S::walk_region(self, ctx, region)
    }

    fn visit_block(&mut self, ctx: &Context, block: Ptr<BasicBlock>) -> WalkResult<B> {
        S::walk_block(self, ctx, block)
    }
}

then the implementation of MyOpAnalyzer becomes:

// Walk tree forward in pre order direction. We return an error value when doing a Break.
impl Visitor<Error> for MyOpAnalyzer {
  fn visit_op(&mut self, ctx: &Context, op: Ptr<Operation>) -> WalkResult<Error> {
    ...
  }
}

visitor5.rs For example removes the Strategy trait in favor of providing walk_x functions. This is more similar to how it is done in rustc_ast. Now the visitors default implementation implement forward iteration only. If users want a different strategy they must overwrite (all) the methods and use one of the given walk functions or implement the iteration themselves.
In visitor5.rs I only implemented a callback function based walker for operations only. The Callback type is somewhat simpler, but the implementation now must take all configurations into account (I'm sure one could have internal special types similar to visitor4.rs for the actual implementations to speed things up an runtime):

impl<B, F> Visitor<B> for Callback<Operation, F>
where
    F: FnMut(&Context, Ptr<Operation>) -> WalkResult<B>,
{
    fn visit_op(&mut self, ctx: &Context, op: Ptr<Operation>) -> WalkResult<B> {
        if self.order == Order::PreOrder {
            let cont = (self.f)(ctx, op)?;
            if cont == WalkCont::Skip {
                // go to next sibling Op, but ignore regions and blocks within this Op
                return WalkResult::Continue(WalkCont::Advance);
            }
        }

        match self.direction {
            Direction::Forward => walk_op(self, ctx, op)?,
            Direction::Reverse => walk_op_rev(self, ctx, op)?,
        };

        if self.order == Order::PostOrder {
            (self.f)(ctx, op)?;
        }
        WalkResult::Continue(WalkCont::Advance)
    }

    fn visit_region(&mut self, ctx: &Context, region: Ptr<Region>) -> WalkResult<B> {
        match self.direction {
            Direction::Forward => walk_region(self, ctx, region),
            Direction::Reverse => walk_region_rev(self, ctx, region),
        }
    }

    fn visit_block(&mut self, ctx: &Context, block: Ptr<BasicBlock>) -> WalkResult<B> {
        match self.direction {
            Direction::Forward => walk_block(self, ctx, block),
            Direction::Reverse => walk_block_rev(self, ctx, block),
        }
    }
}

vs. the type directed approach (which maybe could be optimized better):

impl<B, S: Strategy, F> Visitor<S, B> for Callback<Operation, PreOrder, S, F>
where
    F: FnMut(&Context, Ptr<Operation>) -> WalkResult<B>,
{
    fn visit_op(&mut self, ctx: &Context, op: Ptr<Operation>) -> WalkResult<B> {
        if (self.0)(ctx, op)? == WalkCont::Advance {
            S::walk_op(self, ctx, op)?;
        }
        WalkResult::Continue(WalkCont::Advance)
    }
}

impl<B, S: Strategy, F> Visitor<S, B> for Callback<Operation, PostOrder, S, F>
where
    F: FnMut(&Context, Ptr<Operation>) -> WalkResult<B>,
{
    fn visit_op(&mut self, ctx: &Context, op: Ptr<Operation>) -> WalkResult<B> {
        S::walk_op(self, ctx, op)?;
        (self.0)(ctx, op)
    }
}

[vaivaswatha]

Thank you @urso . visitor5.rs is much simpler and understandable. let me think about it a bit and get back. I'll also check the updated visitor4.rs and get back.

[vaivaswatha]

Closing this as the main topic of discussion, which is a walker is implemented with #32 . The smaller topic of stronger typing for the type system is implemented with #28.