macros: improve decorator handling

unpythonic/syntax/curry.py

@@ -8,6 +8,8 @@
 from macropy.core.hquotes import macros, hq
 from macropy.core.walkers import Walker
 
+from .util import suggest_decorator_index
+
 from ..dynassign import dyn
 from ..fun import curry as curryf, _currycall as currycall
 
@@ -18,8 +20,11 @@ def transform_call(tree, *, stop, **kw):  # technically a node containing the cu
             tree.args = [tree.func] + tree.args
             tree.func = hq[currycall]
         elif type(tree) in (FunctionDef, AsyncFunctionDef):
-            # @curry must run before @trampolined, so put it on the inside
-            tree.decorator_list = tree.decorator_list + [hq[curryf]]
+            k = suggest_decorator_index("curry", tree.decorator_list)
+            if k is not None:
+                tree.decorator_list.insert(k, hq[curryf])
+            else:  # couldn't determine insert position; just plonk it at the end and hope for the best
+                tree.decorator_list.append(hq[curryf])
         elif type(tree) is Lambda:
             # This inserts curry() as the innermost "decorator", and the curry
             # macro is meant to run last (after e.g. tco), so we're fine.

unpythonic/syntax/letdo.py

@@ -172,6 +172,9 @@ def _dletimpl(bindings, fdef, mode, kind):
         values = [t1(rhs) for rhs in values]
     fdef = t2(fdef)
 
+    # We place the let decorator in the innermost position. Hopefully this is ok.
+    # (unpythonic.syntax.util.suggest_decorator_index can't help us here,
+    #  since "let" is not one of the registered decorators)
     letter = dletf if kind == "decorate" else bletf
     bindings = [q[(u[k], ast_literal[v])] for k, v in zip(names, values)]
     fdef.decorator_list = fdef.decorator_list + [hq[letter((ast_literal[bindings],), mode=u[mode], _envname=u[e])]]

unpythonic/syntax/tailtools.py

@@ -22,7 +22,8 @@
 
 from .util import isx, isec, isdo, islet, \
                   detect_callec, detect_lambda, \
-                  has_tco, is_decorator, sort_lambda_decorators
+                  has_tco, is_decorator, sort_lambda_decorators, \
+                  suggest_decorator_index
 from .ifexprs import aif
 from .letdo import let
 
@@ -300,19 +301,12 @@ def _tco_transform_def(tree, *, preproc_cb, **kw):
         if preproc_cb:
             tree = preproc_cb(tree)
         # Enable TCO if not TCO'd already.
-        #
-        # @trampolined needs to be inside of @memoize, otherwise outermost;
-        # so that it is applied **after** any call_ec; this allows also escapes
-        # to return a jump object to the trampoline.
-        # TODO: use unpythonic.regutil.decorator_registry to decide index to insert at
         if not has_tco(tree):
-            ismemoize = [is_decorator(x, "memoize") for x in tree.decorator_list]
-            try:
-                k = ismemoize.index(True) + 1
-                rest = tree.decorator_list[k:] if len(tree.decorator_list) > k else []
-                tree.decorator_list = tree.decorator_list[:k] + [hq[trampolined]] + rest
-            except ValueError:  # no memoize decorator in list
-                tree.decorator_list = [hq[trampolined]] + tree.decorator_list
+            k = suggest_decorator_index("trampolined", tree.decorator_list)
+            if k is not None:
+                tree.decorator_list.insert(k, hq[trampolined])
+            else:  # couldn't determine insert position; just plonk it at the start and hope for the best
+                tree.decorator_list.insert(0, hq[trampolined])
     return tree
 
 # Transform return statements and calls to escape continuations (ec).

unpythonic/syntax/util.py

@@ -35,6 +35,21 @@ def isx(tree, x, accept_attr=True):
            (type(tree) is Captured and tree.name == x) or \
            (accept_attr and type(tree) is Attribute and tree.attr == x)
 
+def getname(tree, accept_attr=True):
+    """The cousin of ``isx``.
+
+    From the same types of trees, extract the name as str.
+
+    If no match on ``tree``, return ``None``.
+    """
+    if type(tree) is Name:
+        return tree.id
+    if type(tree) is Captured:
+        return tree.name
+    if accept_attr and type(tree) is Attribute:
+        return tree.attr
+    return None
+
 def islet(tree, expanded=True):
     """Test whether tree is a ``let[]``, ``letseq[]`` or ``letrec[]``.
 
@@ -272,3 +287,54 @@ def fixit(tree, *, stop, **kw):
             fixit.recurse(thelambda.body)
         return tree
     return fixit.recurse(tree)
+
+# TODO: should we just sort the decorators here, like we do for lambdas?
+# (The current solution is less magic, but less uniform.)
+def suggest_decorator_index(deco_name, decorator_list):
+    """Suggest insertion index for decorator deco_name in given decorator_list.
+
+    The return value ``k`` is intended to be used like this::
+
+        if k is not None:
+            decorator_list.insert(k, mydeco)
+        else:
+            decorator_list.append(mydeco)  # or do something else
+
+    If ``deco_name`` is not in the registry (see ``unpythonic.regutil``),
+    or if an approprite index could not be determined, the return value
+    is ``None``.
+    """
+    if deco_name not in all_decorators:
+        return None  # unknown decorator, don't know where it should go
+    names = [getname(tree) for tree in decorator_list]
+    pri_by_name = {dname: pri for pri, dname in decorator_registry}
+
+    # sanity check that existing known decorators are ordered correctly
+    # (otherwise there is no unique insert position)
+    knownnames = [x for x in names if x in pri_by_name]
+    knownpris = [pri_by_name[x] for x in knownnames]
+    def isascending(lst):
+        maxes = cummax(lst)
+        return all(b >= a for a, b in zip(maxes, maxes[1:]))
+    def cummax(lst):
+        m = float("-inf")
+        out = []
+        for x in lst:
+            m = max(x, m)
+            out.append(m)
+        return out
+    if not (knownpris and isascending(knownpris)):
+        return None
+
+    # when deciding the index, only care about known decorators (hence "suggest")
+    targetpri = pri_by_name[deco_name]
+    if targetpri < knownpris[0]:
+        return 0
+    if targetpri > knownpris[-1]:
+        return len(decorator_list)
+    for pri, dname in zip(knownpris, knownnames):
+        if targetpri >= pri:
+            break
+    else:
+        return None
+    return names.index(dname)