# Source code for ufl.corealg.map_dag

    """Basic algorithms for applying functions to subexpressions."""
    # Copyright (C) 2014-2016 Martin Sandve Alnæs
    #
    # This file is part of UFL (https://www.fenicsproject.org)
    #
    # SPDX-License-Identifier:    LGPL-3.0-or-later
    #
    # Modified by Massimiliano Leoni, 2016

    from ufl.core.expr import Expr
    from ufl.corealg.multifunction import MultiFunction
    from ufl.corealg.traversal import cutoff_unique_post_traversal, unique_post_traversal

    def map_expr_dag(function, expression, compress=True, vcache=None, rcache=None):
        """Apply a function to each subexpression node in an expression DAG.

        If the same function is called multiple times in a transformation
        (as for example in apply_derivatives), then to reuse caches across
        the call, use the arguments vcache and rcache.

        Args:
            function:
                The function
            expression:
                An expression
            compress: If True (default), the output object from
                the function is cached in a dict and reused such that the
                resulting expression DAG does not contain duplicate objects
            vcache: Optional dict for caching results of intermediate
                transformations
            rcache: Optional dict for caching results for compression

        Returns:
            The result of the final function call
        """
        (result,) = map_expr_dags(
            function, [expression], compress=compress, vcache=vcache, rcache=rcache
        )
        return result

    def map_expr_dags(function, expressions, compress=True, vcache=None, rcache=None):
        """Apply a function to each sub-expression node in an expression DAG.

        If *compress* is ``True`` (default) the output object from
        the function is cached in a ``dict`` and reused such that the
        resulting expression DAG does not contain duplicate objects.

        If the same function is called multiple times in a transformation
        (as for example in apply_derivatives), then to reuse caches across
        the call, use the arguments vcache and rcache.

        Args:
            function:
                The function
            expressions:
                An expression
            compress:
                If True (default), the output object from the function is cached in a dict and reused
                such that the resulting expression DAG does not contain duplicate objects
            vcache:
                Optional dict for caching results of intermediate transformations
            rcache:
                Optional dict for caching results for compression

        Returns:
            a list with the result of the final function call for each expression
        """
        # Temporary data structures
        # expr -> r = function(expr,...),  cache of intermediate results
        vcache = {} if vcache is None else vcache
        # r -> r,  cache of result objects for memory reuse
        rcache = {} if rcache is None else rcache

        # Build mapping typecode:bool, for which types to skip the subtree of
        if isinstance(function, MultiFunction):
            cutoff_types = function._is_cutoff_type
            handlers = function._handlers  # Optimization
        else:
            # Regular function: no skipping supported
            cutoff_types = [False] * Expr._ufl_num_typecodes_
            handlers = [function] * Expr._ufl_num_typecodes_

        # Create visited set here to share between traversal calls
        visited = set()

        # Pick faster traversal algorithm if we have no cutoffs
        if any(cutoff_types):

            def traversal(expression):
                return cutoff_unique_post_traversal(expression, cutoff_types, visited)
        else:

            def traversal(expression):
                return unique_post_traversal(expression, visited)

        for expression in expressions:
            # Iterate over all subexpression nodes, child before parent
            for v in traversal(expression):
                # Skip transformations on cache hit
                if v in vcache:
                    continue

                # Cache miss: Get transformed operands, then apply transformation
                if cutoff_types[v._ufl_typecode_]:
                    r = handlers[v._ufl_typecode_](v)
                else:
                    r = handlers[v._ufl_typecode_](v, *(vcache[u] for u in v.ufl_operands))

                # Optionally check if r is in rcache, a memory optimization
                # to be able to keep representation of result compact
                if compress:
                    r2 = rcache.get(r)
                    if r2 is None:
                        # Cache miss: store in rcache
                        rcache[r] = r
                    else:
                        # Cache hit: Use previously computed object r2,
                        # allowing r to be garbage collected as soon as possible
                        r = r2

                # Store result in cache
                vcache[v] = r

        return [vcache[expression] for expression in expressions]