11 meta programming - 01 introduction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Binder](https://mybinder.org/badge_logo.svg)](https://notebooks.gesis.org/binder/v2/gh/jolin-io/fall-in-love-with-julia/main?filepath=11%20meta%20programming%20-%2001%20introduction.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://www.jolin.io\" target=\"_blank\" rel=\"noreferrer noopener\">\n",
    "<img src=\"https://www.jolin.io/assets/Jolin/Jolin-Banner-Website-v1.1-darkmode.webp\">\n",
    "</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fall-in-love-with-Julia: Meta Programming in Julia 101\n",
    "\n",
    "an introduction session\n",
    "\n",
    "I am Stephan Sahm, and today we are going to learn all about meta.\n",
    "\n",
    "1. `Expr` expressions\n",
    "2. `eval` evaluate expressions\n",
    "3. `macro` create syntax helpers\n",
    "4. `@generated` meta functions\n",
    "5. `Cassette.jl` meta programming without `Expr`\n",
    "6. `IRTools.jl` generated functions for IR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# `Expr` expressions\n",
    "\n",
    "Julia is kind of a LISP dialect: Julia's syntax is part of julia itself.\n",
    "Meta programming is builtin."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run some code\n",
    "result = 1 + 2 + 3 + 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# turn it into an expression\n",
    "expr = :(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize expr\n",
    "dump(expr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this is what we want\n",
    "expr = :(1 + 2 + 3 + 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dump(expr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "expr.head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "expr.args"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An alternative to create Expr is using `quote ... end`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "expr = quote\n",
    "    1 + 2 + 3 + 4\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But be careful, that it generates an additional block level,\n",
    "with meta information about where it was constructed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dump(expr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also built `Expr` programmatically"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "expr = Expr(expr.head, expr.args...)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### it's your time\n",
    "\n",
    "👉 try `typeof(expr)` on the `expr` we generated so far"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# your space"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interpolating into Expr\n",
    "\n",
    "`Expr` is a truly flexible data structure: *You can use any arbitary julia value within `Expr`*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "expr = Expr(:call, +, 1, 2, 3, 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dump(expr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "to combine arbitrary values with `Expr` use interpolation via `$`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    ":(2 + result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    ":(2 + $result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "if you want to insert several arguments at once, you can use the splash `...` syntax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    ":(+(1,2,3,4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arguments = (1, 2, 3, 4)\n",
    ":(+($(arguments...)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "same can be done within `quote ... end` blocks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# `eval` evaluate expressions\n",
    "\n",
    "Without the ability to run expressions, they would be pretty useless."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "expr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval(expr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "the macro `@eval` version let's you write normal code, but use expr-interpolation within it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mysymbol = :nice\n",
    "@eval $mysymbol = \"hello world\"\n",
    "nice"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Congratulations! Now you are allmighty 💪🦸 "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### it's your time\n",
    "\n",
    "👉 evaluate a couple of expr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Eval is scoped to a Module\n",
    "\n",
    "So far we evaluated into `Main` module.\n",
    "\n",
    "Note: Also all variables get evaluted as globals into the respective Module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Main.nice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "module MyModule\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MyModule.eval(:(what_do_we_want = \"Climate Justice!\"))\n",
    "MyModule.what_do_we_want"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "the `@eval` macro is centrally defined and works a little different"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@eval MyModule when_do_we_want_it = \"Now!\"\n",
    "MyModule.when_do_we_want_it"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The limitations of eval\n",
    "\n",
    "Julia combines flexibility with performance. Performance means compiling code, however if you can change code all along, how could it be compiled?\n",
    "\n",
    "The answer in julia is **\"world age\"**: Everytime you `eval` a function definition, the world get's older."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Base.get_world_counter()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@eval twice(a) = 2a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Base.get_world_counter()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### it's your time\n",
    "\n",
    "👉 try a couple of `eval` to see what increases the world age and what not"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# your space"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### What does world age mean for you?\n",
    "\n",
    "quote from the paper http://janvitek.org/pubs/oopsla20-j.pdf\n",
    "\n",
    "> Semantically, newly added methods (i.e. ones defined using eval) only become visible when execution returns to the top level,\n",
    "and the set of callable methods for an execution is fixed when it leaves the top level.\n",
    "\n",
    "In my words: Whatever function you create within a function is not visible within the same functioncall."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function test_eval_global_function()\n",
    "    @eval general_advise() = \"give it a second try and it will work\"\n",
    "    general_advise()\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eval_global_function()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function test_eval_global_var()\n",
    "    @eval another_advise = \"if it still doesn't work, change something\"\n",
    "    another_advise\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eval_global_var()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Impressive: global variables can indeed be constructed on the fly.\n",
    "\n",
    "But be careful - this is one of the reasons why global variables are bad for performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function test_eval_global_function_cheat()\n",
    "    @eval cheat() = \"if you really need it, this is how to access a just created function\"\n",
    "    Base.invokelatest(cheat)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eval_global_function_cheat()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# `macro` create syntax helpers\n",
    "\n",
    "\n",
    "`macro = Expr + eval`\n",
    "\n",
    "A macro is a special function which returns an `Expr` that is immediately evaluated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "macro timeit(expr)\n",
    "    quote\n",
    "        before = time()\n",
    "        result = $expr\n",
    "        after = time()\n",
    "        println(\"timedit: $(after - before) seconds\")\n",
    "        result\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@timeit sleep(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### it's your time\n",
    "\n",
    "👉 time a couple of things"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# your space"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Macros under the hood"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@macroexpand @timeit sleep(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is called macro hygiene.\n",
    "\n",
    "New variables defined within the macro are automatically renamed to nonconflicting variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "macro create_variables1()\n",
    "    quote\n",
    "        $(esc(:created_variable1)) = \"works\" \n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@create_variables1\n",
    "created_variable1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "macro create_variables2()\n",
    "    @gensym helper\n",
    "    esc(quote\n",
    "        $helper = \"works too\"\n",
    "        created_variable2 = $helper\n",
    "    end)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@create_variables2\n",
    "created_variable2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@macroexpand @create_variables2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### it's your time\n",
    "\n",
    "👉 inspect what `@show expr` and `@gensym helper` are doing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# your space"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Packages to help you manipulating macros\n",
    "\n",
    "Macros can be quite nasty to work with. There are a couple of packages which can help you:\n",
    "- [MacroTools.jl](https://github.com/FluxML/MacroTools.jl) by FluxML - functional tools\n",
    "- [ExprParsers.jl](https://github.com/jolin-io/ExprParsers.jl) by me - object oriented tools\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# `@generated` meta functions\n",
    "\n",
    "like a macro, but a function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@generated function twice(x)\n",
    "    # Within generated functions, normal printing does not work.\n",
    "    # Think of a generated function as being run at compile-time.\n",
    "    # Luckily there is a more basic alternative:\n",
    "    Core.println(x)\n",
    "    return :(x * x)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "twice(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "twice(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Welcome to the just-in-time version of `eval` 🙂."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### How to inspect generated functions?\n",
    "\n",
    "As of now it is not easily possible to grab the generated code from a generated function.\n",
    "\n",
    "There is even [a stackoverflow question](https://stackoverflow.com/questions/66402105/any-way-to-expand-a-generated-function-in-julia) about it (in julia most often discourse is used instead of stackoverflow)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@code_lowered twice(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Real example: Comparing struct types\n",
    "\n",
    "[StructEquality.jl](https://github.com/jolin-io/StructEquality.jl) is one package of mine which uses @generated functions to solve a simple but common difficulty"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "struct MyArray\n",
    "    value::AbstractArray\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "[\"same\"] == [\"same\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MyArray([\"same\"]) == MyArray([\"same\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "this is unintuitive for most"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using StructEquality: @struct_hash_equal, struct_equal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@struct_hash_equal MyArray\n",
    "\n",
    "MyArray([\"same\"]) == MyArray([\"same\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "much better"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@macroexpand @struct_hash_equal MyArray"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@code_lowered struct_equal(MyArray([\"same\"]), MyArray([\"same\"]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can find the implementation [at github](https://github.com/jolin-io/StructEquality.jl/blob/4756b0906ad0fb742f10aad7c5c017226ee2405a/src/StructEquality.jl#L14-L25) (leave a star if you like it)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# `Cassette.jl` meta programming without `Expr`\n",
    "\n",
    "<p align=\"center\">\n",
    "<img width=\"350px\" src=\"https://raw.githubusercontent.com/JuliaLabs/Cassette.jl/master/docs/img/cassette-logo.png\"/>\n",
    "</p>\n",
    "\n",
    "> Cassette lets you easily extend the Julia language by directly injecting the Julia compiler with new, context-specific behaviors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Cassette\n",
    "Cassette.@context Ctx;\n",
    "Cassette.overdub(Ctx(), /, 1, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This was actually an awesome generated function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@code_lowered 1/2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@code_lowered Cassette.overdub(Ctx(), /, 1, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, `overdub` hooks on the original function definition, but rewrites it\n",
    "such that everything is wrapped into `overdub` and similar helpers. \n",
    "\n",
    "How you use `overdub`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Cassette.@context SinToCosCtx\n",
    "\n",
    "# Override the default recursive `overdub` implementation for `sin(x)`.\n",
    "# Note that there's no tricks here; this is just a normal Julia method\n",
    "# overload using the normal multiple dispatch semantics.\n",
    "Cassette.overdub(::SinToCosCtx, ::typeof(sin), x) = -cos(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = rand(10)\n",
    "y = Cassette.overdub(SinToCosCtx(), sum, i -> cos(i) + sin(i), x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### it's your time\n",
    "\n",
    "👉 build a `overdub` which ignores `println` statements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Cassette.@context PrintlnCtx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# your space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test function\n",
    "function add(a, b)\n",
    "    println(\"I'm about to add $a + $b\")\n",
    "    c = a + b\n",
    "    println(\"c = $c\")\n",
    "    return c\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = rand(3)\n",
    "b = rand(3)\n",
    "add(a, b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Cassette.overdub(PrintlnCtx(), add, a, b)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "👉 extra challenge: collect all the println statements within the context\n",
    "\n",
    "For this you need to know that a Cassette Context can receive a `metadata` keyword."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ctx = PrintlnCtx(metadata = Dict(:key => \"value\"))\n",
    "ctx.metadata[:key]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# your space"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I hope you have seen, how simple it is to use `Cassette.jl`. It is a kind of meta programming, without the use of `Expr`.\n",
    "\n",
    "As `Expr` can easily get messy, I especially like the Cassette approach."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# `IRTools.jl` generated functions for IR\n",
    "\n",
    "IRTools.jl is part of the Flux ecosystem. It is used for computing Gradients in Zygote.jl\n",
    "\n",
    "IR stands for Intermediate Representation. We already saw `@code_lowered` returning `CodeInfo`\n",
    "objects which is such a IR. IRTools.jl offers us an alternative IR, which for some\n",
    "may be easier to work with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function intidentity(a)\n",
    "    s = 0\n",
    "    for i in 1:a\n",
    "        s += 1\n",
    "    end\n",
    "    s\n",
    "end    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@code_lowered intidentity(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using IRTools\n",
    "\n",
    "ir = @code_ir intidentity(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "IRTools allows us to easily access the IR."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "block = IRTools.block(ir, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "IRTools.arguments(block)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "IRTools.branches(block)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "IRTools.predecessors(block)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "IRTools.var(11)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Changing the IR is simple, too."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ir[IRTools.var(11)] = IRTools.xcall(:+, IRTools.var(8), 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, you can make it into a function again"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# first argument to ir is the function type itself, which however is not used here\n",
    "IRTools.evalir(ir, nothing, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "intidentity2 = IRTools.func(ir)\n",
    "intidentity2(nothing, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Cassette.jl can do similar things by working on CodeInfo objects. That is called a contextual `pass` in Cassette.jl\n",
    "See https://julia.mit.edu/Cassette.jl/stable/contextualpass/."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## dynamo\n",
    "\n",
    "`IRTools.@dynamo` allows you to create something like generated functions, but instead of returning an `Expr`, you return an `IRTools.IR`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "IRTools.@dynamo function inspect(a...)\n",
    "    # Within a @dynamo, just like any generated functions, normal printing does not work.\n",
    "    # Think of it as being run at compile-time.\n",
    "    # Luckily there is a more basic alternative:\n",
    "    Core.println(a)\n",
    "    return IRTools.IR(a...)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# note, the first argument is the function itself\n",
    "inspect(prod, 1:4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inspect(prod, 1:4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using MacroTools\n",
    "\n",
    "IRTools.@dynamo function replace_mul_with_sum(a...)\n",
    "    ir = IRTools.IR(a...)\n",
    "    ir = MacroTools.prewalk(ir) do x\n",
    "        x isa GlobalRef && x.name == :* && return GlobalRef(Base, :+)\n",
    "        return x\n",
    "    end\n",
    "    return ir\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "replace_mul_with_sum() do \n",
    "    1 * 2 * 3 * 4\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can even see which IR is returned by your dynamo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@code_ir replace_mul_with_sum() do \n",
    "    1 * 2 * 3 * 4\n",
    "end\n",
    "# @code_lowered works too"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Yet, it does not work recursively"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "replace_mul_with_sum() do \n",
    "    prod(1:4)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But we can make it work by changing our code similar how Cassette does it, where `overdup` always calls `overdub` again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "IRTools.@dynamo function myoverdub(a...)\n",
    "    ir = IRTools.IR(a...)\n",
    "    ir === nothing && return\n",
    "\n",
    "    # explicit version:\n",
    "    for (var, statement) in ir\n",
    "        MacroTools.isexpr(statement.expr, :call) || continue\n",
    "        ir[var] = IRTools.xcall(myoverdub, statement.expr.args...)\n",
    "    end\n",
    "    return ir\n",
    "\n",
    "    # short version:\n",
    "    IRTools.recurse!(ir)\n",
    "    return ir\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "myoverdub(::typeof(*), a, b) = a + b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "myoverdub() do \n",
    "    prod(1:4)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@code_ir myoverdub prod(1:4)  # alternative syntax to see output of dynamo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "🥳 Congratulations 🥳 You've build your own Cassette.jl-like meta programming functionality."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Be careful\n",
    "\n",
    "It is great to know that such tools exist in case you really need it (e.g. it is used for computing automatic gradients in Zygote).\n",
    "\n",
    "BUT BE CAUTIOUS. You can easily destroy everything and changing IR is hard to debug in general.\n",
    "\n",
    "One example: If we would have tried to replace `+` with `*` above (i.e. the other way around), things wouldn't have worked.\n",
    "For one thing, `sum` has a much more complex implementation, which also uses `-`, which will return something totally weird. For another reason, if you would try to write your own `mysum`, you would probably enter an infinite loop. That is because the standard iterator `1:4` will increase its inner state by using `... + 1`, which would translate to `... * 1`, yielding the same state as before and running forever. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Further information\n",
    "- [julia documentation](https://docs.julialang.org/en/v1/manual/metaprogramming/) about macros and generated functions and co\n",
    "- [ExprParsers.jl](https://github.com/jolin-io/ExprParsers.jl) for Expr manipulation\n",
    "- [MacroTools.jl](https://github.com/FluxML/MacroTools.jl) for Expr manipulation\n",
    "- [Cassette.jl docs](https://julia.mit.edu/Cassette.jl/stable/) for further details on Cassette.jl, especially the [compiler pass injection](https://julia.mit.edu/Cassette.jl/stable/contextualpass/) example \n",
    "- [IRTools.jl docs](http://fluxml.ai/IRTools.jl/latest/) for more on intermediate representations. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Thank you for your participation\n",
    "\n",
    "for questions or suggestions please contact me at stephan.sahm@jolin.io\n",
    "\n",
    "\n",
    "#### Sponsored by [Jolin.io](https://www.jolin.io)\n",
    "\n",
    "<a href=\"https://www.jolin.io\" target=\"_blank\" rel=\"noreferrer noopener\">\n",
    "<img src=\"https://www.jolin.io/assets/Jolin/Jolin-Banner-Website-v1.1-darkmode.webp\">\n",
    "</a>\n",
    "\n",
    "Jolin.io is an IT-consultancy focussing on Julia\n",
    "\n",
    "We are there to help you, if you want to\n",
    "- try out Julia at your company, or\n",
    "- transition Matlab, Fortran, R, Python, etc. to Julia\n",
    "- or speed up your existing Julia code"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.8.1",
   "language": "julia",
   "name": "julia-1.8"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.8.1"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}