{
    "cells": [
        {
            "cell_type": "markdown",
            "source": [
                "# Introduction to Redex"
            ],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "\n",
                "## Getting started with combinators\n",
                "\n",
                "Combinators provide a simple and concise way of composing functional code. They may compose, be used by, be mixed with other combinators or standard python functions."
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 1,
            "source": [
                "import operator as op\n",
                "from redex import combinator as cb"
            ],
            "outputs": [],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "The [Serial](https://redex.readthedocs.io/en/latest/_apidoc/redex.combinator/#redex.combinator.serial) combinator allows function composition."
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 2,
            "source": [
                "serial = cb.serial(op.mul, op.add, op.sub)\n",
                "serial(1, 2, 3, 4)"
            ],
            "outputs": [
                {
                    "output_type": "execute_result",
                    "data": {
                        "text/plain": [
                            "1"
                        ]
                    },
                    "metadata": {},
                    "execution_count": 2
                }
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 3,
            "source": [
                "## The same may be achieved with the code:\n",
                "x1 = 1\n",
                "x2 = op.mul(x1, 2)\n",
                "x3 = op.add(x2, 3)\n",
                "x4 = op.sub(x3, 4)\n",
                "x4"
            ],
            "outputs": [
                {
                    "output_type": "execute_result",
                    "data": {
                        "text/plain": [
                            "1"
                        ]
                    },
                    "metadata": {},
                    "execution_count": 3
                }
            ],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "The [Parallel](https://redex.readthedocs.io/en/latest/_apidoc/redex.combinator/#redex.combinator.parallel) combinator applies given functions in parallel to its inputs. Each function consumes a span of inputs. The span sizes are determined by a number of required arguments of these functions."
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 4,
            "source": [
                "parallel = cb.parallel(op.add, op.sub)\n",
                "parallel(1, 2, 3, 4)"
            ],
            "outputs": [
                {
                    "output_type": "execute_result",
                    "data": {
                        "text/plain": [
                            "(3, -1)"
                        ]
                    },
                    "metadata": {},
                    "execution_count": 4
                }
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 5,
            "source": [
                "## The same may be achieved with the code:\n",
                "x1 = op.add(1, 2)\n",
                "x2 = op.sub(3, 4)\n",
                "(x1, x2)"
            ],
            "outputs": [
                {
                    "output_type": "execute_result",
                    "data": {
                        "text/plain": [
                            "(3, -1)"
                        ]
                    },
                    "metadata": {},
                    "execution_count": 5
                }
            ],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "The [Branch](https://redex.readthedocs.io/en/latest/_apidoc/redex.combinator/#redex.combinator.branch) combinator combines multiple branches of given functions and operate on copy of inputs. Each branch includes a sequence of functions applied serially."
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 6,
            "source": [
                "branch = cb.branch(cb.serial(op.mul, op.add), op.sub)\n",
                "branch(1, 2, 3)"
            ],
            "outputs": [
                {
                    "output_type": "execute_result",
                    "data": {
                        "text/plain": [
                            "(5, -1)"
                        ]
                    },
                    "metadata": {},
                    "execution_count": 6
                }
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 7,
            "source": [
                "## The same may be achieved with the code:\n",
                "x1 = op.mul(1, 2)\n",
                "x2 = op.add(x1, 3)\n",
                "x3 = op.sub(1, 2)\n",
                "(x2, x3)"
            ],
            "outputs": [
                {
                    "output_type": "execute_result",
                    "data": {
                        "text/plain": [
                            "(5, -1)"
                        ]
                    },
                    "metadata": {},
                    "execution_count": 7
                }
            ],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "The [Select](https://redex.readthedocs.io/en/latest/_apidoc/redex.combinator/#redex.combinator.select) combinator may be used to rearrange or copy inputs."
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 8,
            "source": [
                "select = cb.select(indices=[0, 0, 1, 1])\n",
                "select(1, 2, 3, 4)"
            ],
            "outputs": [
                {
                    "output_type": "execute_result",
                    "data": {
                        "text/plain": [
                            "(1, 1, 2, 2, 3, 4)"
                        ]
                    },
                    "metadata": {},
                    "execution_count": 8
                }
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 9,
            "source": [
                "## The same may be achieved with the code:\n",
                "xs = (1, 2, 3, 4)\n",
                "(\n",
                "    xs[0],\n",
                "    xs[0],\n",
                "    xs[1],\n",
                "    xs[1],\n",
                "    *xs[2:],\n",
                ")"
            ],
            "outputs": [
                {
                    "output_type": "execute_result",
                    "data": {
                        "text/plain": [
                            "(1, 1, 2, 2, 3, 4)"
                        ]
                    },
                    "metadata": {},
                    "execution_count": 9
                }
            ],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "The [Dup](https://redex.readthedocs.io/en/latest/_apidoc/redex.combinator/#redex.combinator.dup) combinator is a convenient way to make a single copy of inputs."
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 10,
            "source": [
                "dup = cb.dup()\n",
                "dup(1, 2, 3, 4)"
            ],
            "outputs": [
                {
                    "output_type": "execute_result",
                    "data": {
                        "text/plain": [
                            "(1, 1, 2, 3, 4)"
                        ]
                    },
                    "metadata": {},
                    "execution_count": 10
                }
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 11,
            "source": [
                "## The same may be achieved with the code:\n",
                "xs = (1, 2, 3, 4)\n",
                "(\n",
                "    xs[0],\n",
                "    xs[0],\n",
                "    *xs[1:],\n",
                ")"
            ],
            "outputs": [
                {
                    "output_type": "execute_result",
                    "data": {
                        "text/plain": [
                            "(1, 1, 2, 3, 4)"
                        ]
                    },
                    "metadata": {},
                    "execution_count": 11
                }
            ],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "The [Drop](https://redex.readthedocs.io/en/latest/_apidoc/redex.combinator/#redex.combinator.drop) combinator allows to drop some inputs."
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 12,
            "source": [
                "drop = cb.drop(n_in=2)\n",
                "drop(1, 2, 3, 4)"
            ],
            "outputs": [
                {
                    "output_type": "execute_result",
                    "data": {
                        "text/plain": [
                            "(3, 4)"
                        ]
                    },
                    "metadata": {},
                    "execution_count": 12
                }
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 13,
            "source": [
                "## The same may be achieved with the code:\n",
                "xs = (1, 2, 3, 4)\n",
                "xs[2:]"
            ],
            "outputs": [
                {
                    "output_type": "execute_result",
                    "data": {
                        "text/plain": [
                            "(3, 4)"
                        ]
                    },
                    "metadata": {},
                    "execution_count": 13
                }
            ],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "Check out the documentation to find out about [other combinators](https://redex.readthedocs.io/en/latest/_apidoc/redex.combinator)."
            ],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "## Creating new combinators"
            ],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "### Stack\n",
                "\n",
                "Combinators operate on the **stack of function arguments**. They take inputs off the stack, execute a function, then push its outputs back onto the stack. If a function output is a tuple, it gets flattened before placed on the stack. If an input argument is a tuple, each tuple parameter is considered as an independent item on the stack. These parameters are reshaped before get passed to the function as arguments."
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 14,
            "source": [
                "from redex.stack import constrained_call, stackmethod, Stack"
            ],
            "outputs": [],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "The built-in function `operator.add` would take two arguments off the stack and push a single output back onto the stack."
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 15,
            "source": [
                "constrained_call(func=op.add, stack=(1, 2, 0, 0))"
            ],
            "outputs": [
                {
                    "output_type": "execute_result",
                    "data": {
                        "text/plain": [
                            "(3, 0, 0)"
                        ]
                    },
                    "metadata": {},
                    "execution_count": 15
                }
            ],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "The `stackmethod` wrapper provides a convenient way to transform arguments of an arbitrary function to the stack when implementing combinators."
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 16,
            "source": [
                "class Add:\n",
                "    @stackmethod\n",
                "    def __call__(self, stack: Stack) -> Stack:\n",
                "        return constrained_call(func=op.add, stack=stack)\n",
                "\n",
                "add = Add()\n",
                "add(1, 2, 0, 0)"
            ],
            "outputs": [
                {
                    "output_type": "execute_result",
                    "data": {
                        "text/plain": [
                            "(3, 0, 0)"
                        ]
                    },
                    "metadata": {},
                    "execution_count": 16
                }
            ],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "Note that to create a true combinator we would need to derive it from the [Combinator](https://redex.readthedocs.io/en/latest/_apidoc/redex.combinator/#redex.combinator.Combinator) class and specify explicitly the number of its inputs and outputs. In most cases, we would calculate these values dynamically from signatures of nested functions."
            ],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "### Type annotations\n",
                "\n",
                "To pass data through the stack, combinators need to know exact **number of outputs, inputs, and input shapes of functions**. This information may be inferred from type annotations of nested functions or set explicitly.\n",
                "\n",
                "Note that:\n",
                "\n",
                "- when return annotation isn't available, a single output is assumed (to support buit-in functions).\n",
                "- any input argument without default value is counted as a single input.\n",
                "- for tuples used in type annotations, a number of tuple parameters must be definite (e.g. tuple parameters must be specified and variadic tuples must not be used)."
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 17,
            "source": [
                "from redex.function import infer_signature"
            ],
            "outputs": [],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "As an example, the built-in function `operator.add` has two inputs and a single output."
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 18,
            "source": [
                "infer_signature(op.add)"
            ],
            "outputs": [
                {
                    "output_type": "execute_result",
                    "data": {
                        "text/plain": [
                            "Signature(n_in=2, n_out=1, start_index=0, in_shape=((), ()))"
                        ]
                    },
                    "metadata": {},
                    "execution_count": 18
                }
            ],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "Combinators are able to work with functions accepting composite types and returning multiple outputs. To prove that, let's create such functions. The function `create_pair` would accept a single argument of type `int` and output a pair, `tuple[int, int]`. Another function, `add_pairs`, would accept two arguments of the pair type and output a single pair."
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 19,
            "source": [
                "Pair = tuple[int, int]\n",
                "\n",
                "def create_pair(value: int) -> Pair:\n",
                "    return (value, value)\n",
                "\n",
                "def add_pairs(lt: Pair, rt: Pair) -> Pair:\n",
                "    return (lt[0] + rt[0], lt[1] + rt[1])"
            ],
            "outputs": [],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "We see that output `n_out=2` of the function `create_pair` consists of two values. When used with a combinator, `tuple[int, int]` will be flattened and each `int` get placed onto the stack individually."
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 20,
            "source": [
                "infer_signature(create_pair)"
            ],
            "outputs": [
                {
                    "output_type": "execute_result",
                    "data": {
                        "text/plain": [
                            "Signature(n_in=1, n_out=2, start_index=0, in_shape=((),))"
                        ]
                    },
                    "metadata": {},
                    "execution_count": 20
                }
            ],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "The function `add_pairs` accepts two tuples: two parameters per tuple, four inputs in total, `in_n=4`. These four inputs would be taken off the stack, reshaped, and only then passed to the function as arguments."
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 21,
            "source": [
                "infer_signature(add_pairs)"
            ],
            "outputs": [
                {
                    "output_type": "execute_result",
                    "data": {
                        "text/plain": [
                            "Signature(n_in=4, n_out=2, start_index=0, in_shape=(((), ()), ((), ())))"
                        ]
                    },
                    "metadata": {},
                    "execution_count": 21
                }
            ],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "To demonstrate the above functions in action, let's create `create_then_add_pairs` combinator. Given `2` and `3` values, it would create two pairs `(2, 2)` and `(3, 3)` using `create_pair`, then pass these pairs further to `add_pairs`."
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 22,
            "source": [
                "create_then_add_pairs = cb.serial(\n",
                "    cb.parallel(create_pair, create_pair),\n",
                "    add_pairs,\n",
                ")\n",
                "\n",
                "create_then_add_pairs(2, 3)"
            ],
            "outputs": [
                {
                    "output_type": "execute_result",
                    "data": {
                        "text/plain": [
                            "(5, 5)"
                        ]
                    },
                    "metadata": {},
                    "execution_count": 22
                }
            ],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "### Debugging combinators"
            ],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "With debug logging enabled, we can inspect how data flow within combinators.\n",
                "\n"
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 23,
            "source": [
                "import logging"
            ],
            "outputs": [],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "For the above combinator, enabling debug level of the logger gives information about signatures of nested functions and stack usage."
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": 24,
            "source": [
                "logging.getLogger().setLevel(\"DEBUG\")\n",
                "create_then_add_pairs(2, 3)\n",
                "logging.getLogger().setLevel(\"INFO\")"
            ],
            "outputs": [
                {
                    "output_type": "stream",
                    "name": "stderr",
                    "text": [
                        "DEBUG:root:constrained_call :: Parallel             stack_size=2  signature=Signature(n_in=2, n_out=4, start_index=0, in_shape=((), ()))\n",
                        "DEBUG:root:constrained_call :: create_pair          stack_size=1  signature=Signature(n_in=1, n_out=2, start_index=0, in_shape=((),))\n",
                        "DEBUG:root:constrained_call :: create_pair          stack_size=1  signature=Signature(n_in=1, n_out=2, start_index=1, in_shape=((),))\n",
                        "DEBUG:root:constrained_call :: add_pairs            stack_size=4  signature=Signature(n_in=4, n_out=2, start_index=0, in_shape=(((), ()), ((), ())))\n"
                    ]
                }
            ],
            "metadata": {}
        }
    ],
    "metadata": {
        "orig_nbformat": 4,
        "language_info": {
            "name": "python",
            "version": "3.9.6",
            "mimetype": "text/x-python",
            "codemirror_mode": {
                "name": "ipython",
                "version": 3
            },
            "pygments_lexer": "ipython3",
            "nbconvert_exporter": "python",
            "file_extension": ".py"
        },
        "kernelspec": {
            "name": "python3",
            "display_name": "Python 3.9.6 64-bit ('env': venv)"
        },
        "interpreter": {
            "hash": "c565df80b2f03ffa5cfb5208b926034bdb042639b04e4dee7f6bb5d02a0d3735"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 2
}