# Working with the IR

Important

Loki is still under active development and has not yet seen a stable release. Interfaces can change at any time, objects may be renamed, or concepts may be re-thought. Make sure to sync your work to the current release frequently by rebasing feature branches and upstreaming more general applicable work in the form of pull requests.

The most important tool for working with Loki’s internal representation are utilities that traverse the IR to find specific nodes or patterns, to modify or replace subtrees, or to annotate the tree. In Loki there exist two types of tree traversal tools, depending on which level of the IR they operate on:

- Visitors that traverse the tree of control flow nodes;

- Mappers (following Pymbolic’s `pymbolic.mapper` naming convention) that traverse expression trees.

## Visitors

Loki’s visitors work by inspecting the type of each IR node they encounter and then selecting the best matching handler method for that node type. This allows implementing visitors that perform tasks either for very specific node types or generally applicable for any node type, depending on the handler’s name.

Loki includes a range of ready-to-use and configurable visitors for many common use cases, such as discovering certain node types, modifying or replacing nodes in the tree, or creating a string representation of the tree. For some use cases it may be easier to implement new visitors tailored to the task.

### Searching the tree

The first category of visitors traverses the IR and collects a list of results subject to certain criteria. In almost all cases `FindNodes` is the tool for that job with some bespoke variants for specific use cases.

`loki.ir.find.FindNodes`(match[, mode, greedy]) | Find `Node` instances that match a given criterion.
`loki.ir.find.FindScopes`(match[, greedy]) | Find all parent nodes for node `match`.
`loki.ir.find.SequenceFinder`(node_type) | Find repeated nodes of the same type in lists/tuples within a given tree.
`loki.ir.find.PatternFinder`(pattern) | Find a pattern of nodes given as tuple/list of types within a given tree.

A common pattern for using `FindNodes` is the following:

    for loop in FindNodes((Loop, WhileLoop)).visit(routine.body):
        # ...do something with loop...

There are additional visitors that search all expression trees embedded in the control flow IR, which are explained further down.

### Transforming the tree

A core feature of Loki is the ability to transform the IR, which is done using the `Transformer`. It is a visitor that rebuilds the tree and replaces nodes according to a mapper.

`loki.ir.transformer.Transformer`([mapper, ...]) | Visitor class to rebuild the tree and replace nodes according to a mapper.
`loki.ir.transformer.NestedTransformer`([...]) | A `Transformer` that applies replacements in a depth-first fashion.
`loki.ir.transformer.MaskedTransformer`([...]) | An enriched `Transformer` that can selectively include or exclude parts of the tree.
`loki.ir.transformer.NestedMaskedTransformer`([...]) | A `MaskedTransformer` that retains parents for children that are included in the produced tree.

`Transformer` is commonly used in conjunction with `FindNodes`, with the latter being used to build the mapper for the first. The following example removes all loops over the horizontal dimension and replaces them by their body. This code snippet is a simplified version of a transformation used in `ExtractSCATransformation`:

    routine = Subroutine(...)
    horizontal = Dimension(...)

    ...

    loop_map = {}
    for loop in FindNodes(Loop).visit(routine.body):
        if loop.variable == horizontal.variable:
            loop_map[loop] = loop.body
    routine.body = Transformer(loop_map).visit(routine.body)

### Converting the tree to string

The last step in a transformation pipeline is usually to write the transformed IR to a file. This is a task for Loki’s backends which themselves are subclasses of `loki.visitors.pprint.Stringifier`, yet another visitor. `loki.visitors.pprint.Stringifier` doubles as a pretty-printer for the IR that is useful for debugging.

`loki.ir.pprint.Stringifier`([depth, indent, ...]) | Convert a given IR tree to a string representation.
`loki.ir.pprint`(ir[, stream]) | Pretty-print the given IR using `Stringifier`.

### Implementing new visitors

Any new visitor should subclass `Visitor` (or any of its subclasses).

The common base class for all visitors is `GenericVisitor`, declared in `loki.visitors` that provides the basic functionality for matching objects to their handler methods. Derived from that is `Visitor` which adds a default handler `visit_Node` (for `Node`) and functionality to recurse for all items in a list or tuple and return the combined result.

To define handlers in new visitors, they should define `visit_Foo` methods for each class `Foo` they want to handle. If a specific method for a class `Foo` is not found, the MRO of the class is walked in order until a matching method is found (all the way until, for example, `Visitor.visit_Node` applies). The method signature is:

    def visit_Foo(self, o, [*args, **kwargs]):
        pass

The handler is responsible for visiting the children (if any) of the node `o`. `*args` and `**kwargs` may be used to pass information up and down the call stack. You can also pass named keyword arguments, e.g.:

    def visit_Foo(self, o, parent=None, *args, **kwargs):
        pass

## Mappers

Mappers are visitors that traverse expression trees.

They are built upon `pymbolic.mapper` classes and for that reason use a slightly different way of determining the handler methods: each expression tree node (`pymbolic.primitives.Expression`) holds a class attribute `mapper_method` with the name of the relevant method.

Loki provides, similarly to control flow tree visitors, ready-to-use mappers for searching or transforming expression trees, all of which are implemented in `loki.expression.mappers`. In addition, `loki.expression.expr_visitors` provides visitors that apply the same mapper to all expression trees in the IR.

### Searching in expression trees

The equivalent to `FindNodes` for expression trees is `ExpressionRetriever`. Using a generic function handle, (almost) arbitrary conditions can be used as a query that decides whether to include a given node into the list of results.

`loki.expression.mappers.ExpressionRetriever`(query) | A mapper for the expression tree that looks for entries specified by a query.

Note that mappers operate only on expression trees, i.e. using them directly is only useful when working with a single property of a control flow node, such as `loki.ir.Assignment.rhs`. If one wanted to search for expression nodes in all expression trees in the IR, e.g. to find all variables, bespoke visitors exist that apply `ExpressionRetriever` to all expression trees.

For example, the following finds all function calls embedded in expressions (`InlineCall`, as opposed to subroutine calls in `CallStatement`):

    for call in FindInlineCalls().visit(routine.body):
        # ...do something with call...

### Transforming expression trees

Transformations of the expression tree are done very similar to `Transformer`, using the mapper `SubstituteExpressionsMapper` that is given a map to replace matching expression nodes.

`loki.expression.mappers.LokiIdentityMapper`() | A visitor to traverse and transform an expression tree
`loki.expression.mappers.SubstituteExpressionsMapper`(...) | A Pymbolic expression mapper (i.e., a visitor for the expression tree) that defines on-the-fly handlers from a given substitution map.

In the same way that searching can be done on all expression trees in the IR, transformations can be applied to all expression trees at the same time using `SubstituteExpressions`:

The following example shows how expression node discovery and substitution can be combined to replace all occurences of intrinsic function calls. (The code snippet is taken from `replace_intrinsics`, where two dict, `function_map` and `symbol_map`, provide the mapping to rename certain function calls that appear in `routine`.)

    from loki.expression import symbols as sym

    callmap = {}
    for c in FindInlineCalls(unique=False).visit(routine.body):
        cname = c.name.lower()

        if cname in symbol_map:
            callmap[c] = sym.Variable(name=symbol_map[cname], scope=routine.scope)

        if cname in function_map:
            fct_symbol = sym.ProcedureSymbol(function_map[cname], scope=routine.scope)
            callmap[c] = sym.InlineCall(fct_symbol, parameters=c.parameters,
                                        kw_parameters=c.kw_parameters)

    routine.body = SubstituteExpressions(callmap).visit(routine.body)

### Converting expressions to string

Every backend has their own mapper to convert expressions to a source code string, according to the corresponding language specification. All build on a common base class `LokiStringifyMapper`, which is also called automatically when converting any expression node to string.

`loki.expression.mappers.LokiStringifyMapper`(...) | A class derived from the default `StringifyMapper` that adds mappings for nodes of the expression tree that we added ourselves.