[pypy-commit] pypy online-transforms: make it possible to register transformers

[rlamy]

Author: Ronan Lamy <ronan.lamy at gmail.com>
Branch: online-transforms
Changeset: r73970:39658f95ae7f
Date: 2014-10-15 21:59 +0100
http://bitbucket.org/pypy/pypy/changeset/39658f95ae7f/

Log:	make it possible to register transformers

diff --git a/rpython/annotator/unaryop.py b/rpython/annotator/unaryop.py
--- a/rpython/annotator/unaryop.py
+++ b/rpython/annotator/unaryop.py
@@ -5,6 +5,7 @@
 from __future__ import absolute_import
 
 from rpython.flowspace.operation import op
+from rpython.flowspace.model import const
 from rpython.annotator.model import (SomeObject, SomeInteger, SomeBool,
     SomeString, SomeChar, SomeList, SomeDict, SomeTuple, SomeImpossibleValue,
     SomeUnicodeCodePoint, SomeInstance, SomeBuiltin, SomeBuiltinMethod,
@@ -705,6 +706,12 @@
     def setslice(self, s_start, s_stop, s_iterable):
         return self._emulate_call('__setslice__', s_start, s_stop, s_iterable)
 
+ at op.len.register_transform(SomeInstance)
+def len_SomeInstance(annotator, v_arg):
+    get_len = op.getattr(v_arg, const('__len__'))
+    return [get_len, op.simple_call(get_len.result)]
+
+
 class __extend__(SomeBuiltin):
     def call(self, args, implicit_init=False):
         args_s, kwds = args.unpack()
diff --git a/rpython/flowspace/operation.py b/rpython/flowspace/operation.py
--- a/rpython/flowspace/operation.py
+++ b/rpython/flowspace/operation.py
@@ -57,8 +57,10 @@
             setattr(op, cls.opname, cls)
         if cls.dispatch == 1:
             cls._registry = {}
+            cls._transform = {}
         elif cls.dispatch == 2:
             cls._registry = DoubleDispatchRegistry()
+            cls._transform = DoubleDispatchRegistry()
 
 
 class HLOperation(SpaceOperation):
@@ -104,8 +106,13 @@
     def get_can_only_throw(self, annotator):
         return None
 
+    def get_transformer(self, *args_s):
+        return lambda *args: None
+
     def transform(self, annotator):
-        pass
+        args_s = [annotator.annotation(arg) for arg in self.args]
+        transformer = self.get_transformer(*args_s)
+        return transformer(annotator, *self.args)
 
 class PureOperation(HLOperation):
     pure = True
@@ -188,6 +195,37 @@
         except AttributeError:
             return cls._dispatch(type(s_arg))
 
+    @classmethod
+    def get_specialization(cls, s_arg, *_ignored):
+        try:
+            impl = getattr(s_arg, cls.opname)
+
+            def specialized(annotator, arg, *other_args):
+                return impl(*[annotator.annotation(x) for x in other_args])
+            try:
+                specialized.can_only_throw = impl.can_only_throw
+            except AttributeError:
+                pass
+            return specialized
+        except AttributeError:
+            return cls._dispatch(type(s_arg))
+
+    @classmethod
+    def register_transform(cls, Some_cls):
+        def decorator(func):
+            cls._transform[Some_cls] = func
+            return func
+        return decorator
+
+    @classmethod
+    def get_transformer(cls, s_arg, *_ignored):
+        for c in type(s_arg).__mro__:
+            try:
+                return cls._transform[c]
+            except KeyError:
+                pass
+        return lambda *args: None
+
 
 class DoubleDispatchMixin(object):
     dispatch = 2
@@ -219,6 +257,20 @@
         spec = type(self).get_specialization(*args_s)
         return read_can_only_throw(spec, args_s[0], args_s[1])
 
+    @classmethod
+    def register_transform(cls, Some1, Some2):
+        def decorator(func):
+            cls._transform[Some1, Some2] = func
+            return func
+        return decorator
+
+    @classmethod
+    def get_transformer(cls, s_arg1, s_arg2, *_ignored):
+        try:
+            return cls._transform[type(s_arg1), type(s_arg2)]
+        except KeyError:
+            return lambda *args: None
+
 
 def add_operator(name, arity, dispatch=None, pyfunc=None, pure=False, ovf=False):
     operator_func = getattr(operator, name, None)
@@ -606,15 +658,6 @@
         return ArgumentsForTranslation.fromshape(args_s[0].const,
                                                 list(args_s[1:]))
 
-def transform_len(hlop, annotator):
-    from rpython.annotator.model import SomeInstance
-    s_arg = annotator.annotation(hlop.args[0])
-    if isinstance(s_arg, SomeInstance):
-        get_len = op.getattr(hlop.args[0], const('__len__'))
-        return [get_len, op.simple_call(get_len.result)]
-
-op.len.transform = transform_len
-
 
 # Other functions that get directly translated to SpaceOperators
 func2op[type] = op.type