Optimize setitem() by calling mp_ass_subscript if available and it's not the generic slot impl

use it for attrwrapper_cls (This gets also often called for real dicts but the implement the mapping already)

src/capi/typeobject.cpp

@@ -1316,7 +1316,7 @@ static int method_is_overloaded(PyObject* left, PyObject* right, const char* nam
 
 SLOT1(slot_mp_subscript, "__getitem__", PyObject*, "O")
 
-static int slot_mp_ass_subscript(PyObject* self, PyObject* key, PyObject* value) noexcept {
+int slot_mp_ass_subscript(PyObject* self, PyObject* key, PyObject* value) noexcept {
     STAT_TIMER(t0, "us_timer_slot_mpasssubscript", SLOT_AVOIDABILITY(self));
 
     PyObject* res;

src/capi/typeobject.h

@@ -47,6 +47,7 @@ PyObject* slot_tp_iter(PyObject* self) noexcept;
 PyObject* slot_tp_iternext(PyObject* self) noexcept;
 PyObject* slot_tp_new(PyTypeObject* self, PyObject* args, PyObject* kwds) noexcept;
 PyObject* slot_mp_subscript(PyObject* self, PyObject* arg1) noexcept;
+int slot_mp_ass_subscript(PyObject* self, PyObject* key, PyObject* value) noexcept;
 int slot_sq_contains(PyObject* self, PyObject* value) noexcept;
 Py_ssize_t slot_sq_length(PyObject* self) noexcept;
 PyObject* slot_tp_getattr_hook(PyObject* self, PyObject* name) noexcept;

src/core/types.h

@@ -482,7 +482,7 @@ inline BoxedString* internStringMortal(llvm::StringRef s) {
 }
 
 // TODO this is an immortal intern for now
-inline void internStringMortalInplace(BoxedString*& s) {
+inline void internStringMortalInplace(BoxedString*& s) noexcept {
     PyString_InternInPlace((PyObject**)&s);
 }
 

src/runtime/builtin_modules/builtins.cpp

@@ -464,7 +464,7 @@ Box* bltinImport(Box* name, Box* globals, Box* locals, Box** args) {
     // Well, it gets passed to PyImport_ImportModuleLevel() and then import_module_level(),
     // which ignores it.  So we don't even pass it through.
 
-    name = coerceUnicodeToStr(name);
+    name = coerceUnicodeToStr<CXX>(name);
 
     if (name->cls != str_cls) {
         raiseExcHelper(TypeError, "__import__() argument 1 must be string, not %s", getTypeName(name));
@@ -479,7 +479,7 @@ Box* bltinImport(Box* name, Box* globals, Box* locals, Box** args) {
 }
 
 Box* delattrFunc(Box* obj, Box* _str) {
-    _str = coerceUnicodeToStr(_str);
+    _str = coerceUnicodeToStr<CXX>(_str);
 
     if (_str->cls != str_cls)
         raiseExcHelper(TypeError, "attribute name must be string, not '%s'", getTypeName(_str));
@@ -543,15 +543,9 @@ Box* getattrFuncInternal(BoxedFunctionBase* func, CallRewriteArgs* rewrite_args,
         }
     }
 
-    try {
-        _str = coerceUnicodeToStr(_str);
-    } catch (ExcInfo e) {
-        if (S == CAPI) {
-            setCAPIException(e);
-            return NULL;
-        } else
-            throw e;
-    }
+    _str = coerceUnicodeToStr<S>(_str);
+    if (S == CAPI && !_str)
+        return NULL;
 
     if (!PyString_Check(_str)) {
         if (S == CAPI) {
@@ -601,7 +595,7 @@ Box* getattrFuncInternal(BoxedFunctionBase* func, CallRewriteArgs* rewrite_args,
 }
 
 Box* setattrFunc(Box* obj, Box* _str, Box* value) {
-    _str = coerceUnicodeToStr(_str);
+    _str = coerceUnicodeToStr<CXX>(_str);
 
     if (_str->cls != str_cls) {
         raiseExcHelper(TypeError, "setattr(): attribute name must be string");
@@ -651,15 +645,9 @@ Box* hasattrFuncInternal(BoxedFunctionBase* func, CallRewriteArgs* rewrite_args,
             rewrite_args->arg2->addGuard((intptr_t)arg2);
     }
 
-    try {
-        _str = coerceUnicodeToStr(_str);
-    } catch (ExcInfo e) {
-        if (S == CAPI) {
-            setCAPIException(e);
-            return NULL;
-        } else
-            throw e;
-    }
+    _str = coerceUnicodeToStr<S>(_str);
+    if (S == CAPI && !_str)
+        return NULL;
 
     if (!PyString_Check(_str)) {
         if (S == CAPI) {

src/runtime/dict.cpp

@@ -432,6 +432,7 @@ static int dict_ass_sub(PyDictObject* mp, PyObject* v, PyObject* w) noexcept {
         assert(res == None);
     } catch (ExcInfo e) {
         setCAPIException(e);
+        return -1;
     }
     return 0;
 }

src/runtime/import.cpp

@@ -683,7 +683,7 @@ extern "C" Box* import(int level, Box* from_imports, llvm::StringRef module_name
 }
 
 Box* impFindModule(Box* _name, BoxedList* path) {
-    _name = coerceUnicodeToStr(_name);
+    _name = coerceUnicodeToStr<CXX>(_name);
     RELEASE_ASSERT(_name->cls == str_cls, "");
 
     BoxedString* name = static_cast<BoxedString*>(_name);

src/runtime/objmodel.cpp

@@ -3488,7 +3488,7 @@ void rearrangeArguments(ParamReceiveSpec paramspec, const ParamNames* param_name
         }
 
         for (const auto& p : *d_kwargs) {
-            auto k = coerceUnicodeToStr(p.first);
+            auto k = coerceUnicodeToStr<CXX>(p.first);
 
             if (k->cls != str_cls)
                 raiseExcHelper(TypeError, "%s() keywords must be strings", func_name);
@@ -4720,7 +4720,6 @@ static Box* callItemAttr(Box* target, BoxedString* item_str, Box* item, Box* val
 template <ExceptionStyle S>
 static Box* callItemOrSliceAttr(Box* target, BoxedString* item_str, BoxedString* slice_str, Box* slice, Box* value,
                                 CallRewriteArgs* rewrite_args) noexcept(S == CAPI) {
-
     // This function contains a lot of logic for deciding between whether to call
     // the slice operator or the item operator, so we can match CPython's behavior
     // on custom classes that define those operators. However, for builtin types,
@@ -4979,6 +4978,12 @@ extern "C" Box* getitem_capi(Box* target, Box* slice) noexcept {
     return rtn;
 }
 
+static void setitemHelper(Box* target, Box* slice, Box* value) {
+    int ret = target->cls->tp_as_mapping->mp_ass_subscript(target, slice, value);
+    if (ret == -1)
+        throwCAPIException();
+}
+
 // target[slice] = value
 extern "C" void setitem(Box* target, Box* slice, Box* value) {
     STAT_TIMER(t0, "us_timer_slowpath_setitem", 10);
@@ -4992,6 +4997,23 @@ extern "C" void setitem(Box* target, Box* slice, Box* value) {
     static BoxedString* setitem_str = internStringImmortal("__setitem__");
     static BoxedString* setslice_str = internStringImmortal("__setslice__");
 
+    auto&& m = target->cls->tp_as_mapping;
+    if (m && m->mp_ass_subscript && m->mp_ass_subscript != slot_mp_ass_subscript) {
+        if (rewriter.get()) {
+            RewriterVar* r_obj = rewriter->getArg(0);
+            RewriterVar* r_slice = rewriter->getArg(1);
+            RewriterVar* r_value = rewriter->getArg(2);
+            RewriterVar* r_cls = r_obj->getAttr(offsetof(Box, cls));
+            RewriterVar* r_m = r_cls->getAttr(offsetof(BoxedClass, tp_as_mapping));
+            r_m->addGuardNotEq(0);
+            rewriter->call(true, (void*)setitemHelper, r_obj, r_slice, r_value);
+            rewriter->commit();
+        }
+
+        setitemHelper(target, slice, value);
+        return;
+    }
+
     Box* rtn;
     if (rewriter.get()) {
         CallRewriteArgs rewrite_args(rewriter.get(), rewriter->getArg(0), rewriter->getReturnDestination());
@@ -5502,7 +5524,7 @@ Box* _typeNew(BoxedClass* metatype, BoxedString* name, BoxedTuple* bases, BoxedD
     }
 
     for (const auto& p : *attr_dict) {
-        auto k = coerceUnicodeToStr(p.first);
+        auto k = coerceUnicodeToStr<CXX>(p.first);
 
         RELEASE_ASSERT(k->cls == str_cls, "");
         BoxedString* s = static_cast<BoxedString*>(k);
@@ -5831,7 +5853,7 @@ extern "C" Box* importStar(Box* _from_module, Box* to_globals) {
             }
             idx++;
 
-            attr_name = coerceUnicodeToStr(attr_name);
+            attr_name = coerceUnicodeToStr<CXX>(attr_name);
 
             if (attr_name->cls != str_cls)
                 raiseExcHelper(TypeError, "attribute name must be string, not '%s'", getTypeName(attr_name));

src/runtime/types.cpp

@@ -2246,7 +2246,7 @@ public:
         RELEASE_ASSERT(_self->cls == attrwrapper_cls, "");
         AttrWrapper* self = static_cast<AttrWrapper*>(_self);
 
-        _key = coerceUnicodeToStr(_key);
+        _key = coerceUnicodeToStr<CXX>(_key);
 
         RELEASE_ASSERT(_key->cls == str_cls, "");
         BoxedString* key = static_cast<BoxedString*>(_key);
@@ -2256,11 +2256,27 @@ public:
         return None;
     }
 
+    static int ass_sub(PyDictObject* mp, PyObject* v, PyObject* w) noexcept {
+        try {
+            Box* res;
+            if (w == NULL) {
+                res = AttrWrapper::delitem((Box*)mp, v);
+            } else {
+                res = AttrWrapper::setitem((Box*)mp, v, w);
+            }
+            assert(res == None);
+        } catch (ExcInfo e) {
+            setCAPIException(e);
+            return -1;
+        }
+        return 0;
+    }
+
     static Box* setdefault(Box* _self, Box* _key, Box* value) {
         RELEASE_ASSERT(_self->cls == attrwrapper_cls, "");
         AttrWrapper* self = static_cast<AttrWrapper*>(_self);
 
-        _key = coerceUnicodeToStr(_key);
+        _key = coerceUnicodeToStr<CXX>(_key);
 
         RELEASE_ASSERT(_key->cls == str_cls, "");
         BoxedString* key = static_cast<BoxedString*>(_key);
@@ -2277,7 +2293,7 @@ public:
         RELEASE_ASSERT(_self->cls == attrwrapper_cls, "");
         AttrWrapper* self = static_cast<AttrWrapper*>(_self);
 
-        _key = coerceUnicodeToStr(_key);
+        _key = coerceUnicodeToStr<CXX>(_key);
 
         RELEASE_ASSERT(_key->cls == str_cls, "");
         BoxedString* key = static_cast<BoxedString*>(_key);
@@ -2293,7 +2309,9 @@ public:
         RELEASE_ASSERT(_self->cls == attrwrapper_cls, "");
         AttrWrapper* self = static_cast<AttrWrapper*>(_self);
 
-        _key = coerceUnicodeToStr(_key);
+        _key = coerceUnicodeToStr<S>(_key);
+        if (S == CAPI && !_key)
+            return NULL;
 
         RELEASE_ASSERT(_key->cls == str_cls, "%s", _key->cls->tp_name);
         BoxedString* key = static_cast<BoxedString*>(_key);
@@ -2313,7 +2331,7 @@ public:
         RELEASE_ASSERT(_self->cls == attrwrapper_cls, "");
         AttrWrapper* self = static_cast<AttrWrapper*>(_self);
 
-        _key = coerceUnicodeToStr(_key);
+        _key = coerceUnicodeToStr<CXX>(_key);
 
         RELEASE_ASSERT(_key->cls == str_cls, "");
         BoxedString* key = static_cast<BoxedString*>(_key);
@@ -2334,7 +2352,7 @@ public:
         RELEASE_ASSERT(_self->cls == attrwrapper_cls, "");
         AttrWrapper* self = static_cast<AttrWrapper*>(_self);
 
-        _key = coerceUnicodeToStr(_key);
+        _key = coerceUnicodeToStr<CXX>(_key);
 
         RELEASE_ASSERT(_key->cls == str_cls, "%s", _key->cls->tp_name);
         BoxedString* key = static_cast<BoxedString*>(_key);
@@ -2371,11 +2389,13 @@ public:
         return boxString(os.str());
     }
 
-    static Box* contains(Box* _self, Box* _key) {
+    template <ExceptionStyle S> static Box* contains(Box* _self, Box* _key) noexcept(S == CAPI) {
         RELEASE_ASSERT(_self->cls == attrwrapper_cls, "");
         AttrWrapper* self = static_cast<AttrWrapper*>(_self);
 
-        _key = coerceUnicodeToStr(_key);
+        _key = coerceUnicodeToStr<S>(_key);
+        if (S == CAPI && !_key)
+            return NULL;
 
         RELEASE_ASSERT(_key->cls == str_cls, "");
         BoxedString* key = static_cast<BoxedString*>(_key);
@@ -2385,6 +2405,13 @@ public:
         return r ? True : False;
     }
 
+    static int sq_contains(Box* _self, Box* _key) noexcept {
+        Box* rtn = contains<CAPI>(_self, _key);
+        if (!rtn)
+            return -1;
+        return rtn == True;
+    }
+
     static Box* keys(Box* _self) {
         RELEASE_ASSERT(_self->cls == attrwrapper_cls, "");
         AttrWrapper* self = static_cast<AttrWrapper*>(_self);
@@ -2706,7 +2733,7 @@ Box* objectHash(Box* obj) {
 }
 
 Box* objectSetattr(Box* obj, Box* attr, Box* value) {
-    attr = coerceUnicodeToStr(attr);
+    attr = coerceUnicodeToStr<CXX>(attr);
     if (attr->cls != str_cls) {
         raiseExcHelper(TypeError, "attribute name must be string, not '%s'", attr->cls->tp_name);
     }
@@ -3843,6 +3870,8 @@ void setupRuntime() {
 
     static PyMappingMethods attrwrapper_as_mapping;
     attrwrapper_cls->tp_as_mapping = &attrwrapper_as_mapping;
+    static PySequenceMethods attrwrapper_as_sequence;
+    attrwrapper_cls->tp_as_sequence = &attrwrapper_as_sequence;
     attrwrapper_cls->giveAttr("__setitem__", new BoxedFunction(boxRTFunction((void*)AttrWrapper::setitem, UNKNOWN, 3)));
     attrwrapper_cls->giveAttr(
         "pop", new BoxedFunction(boxRTFunction((void*)AttrWrapper::pop, UNKNOWN, 3, false, false), { NULL }));
@@ -3855,7 +3884,7 @@ void setupRuntime() {
         "get", new BoxedFunction(boxRTFunction((void*)AttrWrapper::get, UNKNOWN, 3, false, false), { None }));
     attrwrapper_cls->giveAttr("__str__", new BoxedFunction(boxRTFunction((void*)AttrWrapper::str, UNKNOWN, 1)));
     attrwrapper_cls->giveAttr("__contains__",
-                              new BoxedFunction(boxRTFunction((void*)AttrWrapper::contains, UNKNOWN, 2)));
+                              new BoxedFunction(boxRTFunction((void*)AttrWrapper::contains<CXX>, UNKNOWN, 2)));
     attrwrapper_cls->giveAttr("__eq__", new BoxedFunction(boxRTFunction((void*)AttrWrapper::eq, UNKNOWN, 2)));
     attrwrapper_cls->giveAttr("__ne__", new BoxedFunction(boxRTFunction((void*)AttrWrapper::ne, UNKNOWN, 2)));
     attrwrapper_cls->giveAttr("keys", new BoxedFunction(boxRTFunction((void*)AttrWrapper::keys, LIST, 1)));
@@ -3874,6 +3903,8 @@ void setupRuntime() {
     attrwrapper_cls->freeze();
     attrwrapper_cls->tp_iter = AttrWrapper::iter;
     attrwrapper_cls->tp_as_mapping->mp_subscript = (binaryfunc)AttrWrapper::getitem<CAPI>;
+    attrwrapper_cls->tp_as_mapping->mp_ass_subscript = (objobjargproc)AttrWrapper::ass_sub;
+    attrwrapper_cls->tp_as_sequence->sq_contains = (objobjproc)AttrWrapper::sq_contains;
 
     attrwrapperiter_cls->giveAttr("__hasnext__",
                                   new BoxedFunction(boxRTFunction((void*)AttrWrapperIter::hasnext, UNKNOWN, 1)));

src/runtime/util.cpp

@@ -104,7 +104,7 @@ Box* noneIfNull(Box* b) {
     }
 }
 
-Box* coerceUnicodeToStr(Box* unicode) {
+template <ExceptionStyle S> Box* coerceUnicodeToStr(Box* unicode) noexcept(S == CAPI) {
     if (!isSubclass(unicode->cls, unicode_cls))
         return unicode;
 
@@ -117,6 +117,10 @@ Box* coerceUnicodeToStr(Box* unicode) {
     return r;
 }
 
+// force instantiation:
+template Box* coerceUnicodeToStr<CXX>(Box* unicode);
+template Box* coerceUnicodeToStr<CAPI>(Box* unicode);
+
 Box* boxStringFromCharPtr(const char* s) {
     return boxString(s);
 }

src/runtime/util.h

@@ -56,7 +56,7 @@ Box* boxStringFromCharPtr(const char* s);
 // strings and a string that happens to be its encoding.  It seems safer to just encode as ascii,
 // which will throw an exception if you try to pass something that might run into this risk.
 // (We wrap the unicode error and throw a TypeError)
-Box* coerceUnicodeToStr(Box* unicode);
+template <ExceptionStyle S> Box* coerceUnicodeToStr(Box* unicode) noexcept(S == CAPI);
 
 extern "C" bool hasnext(Box* o);
 extern "C" void dump(void* p);