Fix adding methods to existing classes

basis/cocoa/cocoa-tests.factor

@@ -1,7 +1,7 @@
 IN: cocoa.tests
 USING: cocoa cocoa.messages cocoa.subclassing cocoa.types
 compiler kernel namespaces cocoa.classes tools.test memory
-compiler.units ;
+compiler.units math ;
 
 CLASS: {
     { +superclass+ "NSObject" }
@@ -45,3 +45,27 @@ Bar [
 [ 2.0 ] [ "x" get NSRect-y ] unit-test
 [ 101.0 ] [ "x" get NSRect-w ] unit-test
 [ 102.0 ] [ "x" get NSRect-h ] unit-test
+
+! Make sure that we can add methods
+CLASS: {
+    { +superclass+ "NSObject" }
+    { +name+ "Bar" }
+} {
+    "bar"
+    "NSRect"
+    { "id" "SEL" }
+    [ 2drop test-foo "x" get ]
+} {
+    "babb"
+    "int"
+    { "id" "SEL" "int" }
+    [ 2nip sq ]
+} ;
+
+[ 144 ] [
+    Bar [
+        -> alloc -> init
+        dup 12 -> babb
+        swap -> release
+    ] compile-call
+] unit-test

basis/cocoa/subclassing/subclassing.factor

@@ -1,10 +1,9 @@
-! Copyright (C) 2006, 2008 Slava Pestov
+! Copyright (C) 2006, 2008 Slava Pestov, Joe Groff.
 ! See http://factorcode.org/license.txt for BSD license.
 USING: alien alien.c-types alien.strings arrays assocs
 combinators compiler hashtables kernel libc math namespaces
-parser sequences words cocoa.messages cocoa.runtime
+parser sequences words cocoa.messages cocoa.runtime locals
-compiler.units io.encodings.ascii generalizations
+compiler.units io.encodings.ascii continuations make fry ;
-continuations make ;
 IN: cocoa.subclassing
 
 : init-method ( method -- sel imp types )
@@ -12,22 +11,25 @@ IN: cocoa.subclassing
     [ sel_registerName ] [ execute ] [ ascii string>alien ]
     tri* ;
 
-: throw-if-false ( YES/NO -- )
+: throw-if-false ( obj what -- )
-    zero? [ "Failed to add method or protocol to class" throw ]
+    swap { f 0 } member?
-    when ;
+    [ "Failed to " prepend throw ] [ drop ] if ;
+
+: add-method ( class sel imp types -- )
+    class_addMethod "add method to class" throw-if-false ;
 
 : add-methods ( methods class -- )
-    swap
+    '[ [ _ ] dip init-method add-method ] each ;
-    [ init-method class_addMethod throw-if-false ] with each ;
+
+: add-protocol ( class protocol -- )
+    class_addProtocol "add protocol to class" throw-if-false ;
 
 : add-protocols ( protocols class -- )
-    swap [ objc-protocol class_addProtocol throw-if-false ]
+    '[ [ _ ] dip objc-protocol add-protocol ] each ;
-    with each ;
 
-: (define-objc-class) ( protocols superclass name imeth -- )
+: (define-objc-class) ( imeth protocols superclass name -- )
-    -rot
     [ objc-class ] dip 0 objc_allocateClassPair
-    [ add-methods ] [ add-protocols ] [ objc_registerClassPair ]
+    [ add-protocols ] [ add-methods ] [ objc_registerClassPair ]
     tri ;
 
 : encode-types ( return types -- encoding )
@@ -45,28 +47,19 @@ IN: cocoa.subclassing
         [ first4 prepare-method 3array ] map
     ] with-compilation-unit ;
 
-: types= ( a b -- ? )
+:: (redefine-objc-method) ( class method -- )
-    [ ascii alien>string ] bi@ = ;
+    method init-method [| sel imp types |
-
+        class sel class_getInstanceMethod [
-: (verify-method-type) ( class sel types -- )
+            imp method_setImplementation drop
-    [ class_getInstanceMethod method_getTypeEncoding ]
+        ] [
-    dip types=
+            class sel imp types add-method
-    [ "Objective-C method types cannot be changed once defined" throw ]
+        ] if*
-    unless ;
+    ] call ;
-: verify-method-type ( class sel imp types -- class sel imp types )
-    4 ndup nip (verify-method-type) ;
-
-: (redefine-objc-method) ( class method -- )
-    init-method ! verify-method-type
-    drop
-    [ class_getInstanceMethod ] dip method_setImplementation drop ;
     
 : redefine-objc-methods ( imeth name -- )
     dup class-exists? [
-        objc_getClass swap [ (redefine-objc-method) ] with each
+        objc_getClass '[ [ _ ] dip (redefine-objc-method) ] each
-    ] [
+    ] [ 2drop ] if ;
-        2drop
-    ] if ;
 
 SYMBOL: +name+
 SYMBOL: +protocols+
@@ -76,10 +69,10 @@ SYMBOL: +superclass+
     clone [
         prepare-methods
         +name+ get "cocoa.classes" create drop
-        +name+ get 2dup redefine-objc-methods swap [
+        +name+ get 2dup redefine-objc-methods swap
-            +protocols+ get , +superclass+ get , +name+ get , ,
+        +protocols+ get +superclass+ get +name+ get
-            \ (define-objc-class) ,
+        '[ _ _ _ _ (define-objc-class) ]
-        ] [ ] make import-objc-class
+        import-objc-class
     ] bind ;
 
 : CLASS: