diff --git a/src/reagent/core.cljs b/src/reagent/core.cljs
index 0ce45e9..bad5cd6 100644
--- a/src/reagent/core.cljs
+++ b/src/reagent/core.cljs
@@ -71,12 +71,14 @@ Everything is optional, except :render.
   "Set state of a component."
   [this new-state]
   (assert (util/reagent-component? this))
+  (assert (or (nil? new-state) (map? new-state)))
   (comp/replace-state this new-state))
 
 (defn set-state
   "Merge component state with new-state."
   [this new-state]
   (assert (util/reagent-component? this))
+  (assert (or (nil? new-state) (map? new-state)))
   (comp/set-state this new-state))
 
 
@@ -103,10 +105,10 @@ Everything is optional, except :render.
   [this]
   (.getDOMNode this))
 
-(defn ref [parent ref child]
+(defn ref [parent key child]
   (assert (util/reagent-component? parent))
-  (assert (keyword? ref) (str "Ref must be a keyword, not " (pr-str ref)))
-  (tmpl/make-ref-component parent ref child))
+  (assert (keyword? key) (str "Key must be a keyword, not " (pr-str key)))
+  (tmpl/make-ref-component parent key child))
 
 (defn refs [this]
   (assert (util/reagent-component? this))
diff --git a/src/reagent/impl/component.cljs b/src/reagent/impl/component.cljs
index 425a04d..4e677da 100644
--- a/src/reagent/impl/component.cljs
+++ b/src/reagent/impl/component.cljs
@@ -12,19 +12,22 @@
 
 ;;; State
 
+(defn state-atom [this]
+  (let [sa (aget this cljs-state)]
+    (if-not (nil? sa)
+      sa
+      (aset this cljs-state (ratom/atom nil)))))
+
 (defn state [this]
-  (aget this cljs-state))
+  (deref (state-atom this)))
 
 (defn replace-state [this new-state]
   ;; Don't use React's replaceState, since it doesn't play well
   ;; with clojure maps
-  (let [old-state (state this)]
-    (when-not (identical? old-state new-state)
-      (aset this cljs-state new-state)
-      (.forceUpdate this))))
+  (reset! (state-atom this) new-state))
 
 (defn set-state [this new-state]
-  (replace-state this (merge (state this) new-state)))
+  (swap! (state-atom this) merge new-state))
 
 
 ;;; Rendering
@@ -64,8 +67,7 @@
     :getInitialState
     (fn []
       (this-as C
-               (when f
-                 (aset C cljs-state (merge (state C) (f C))))))
+               (set-state C (f C))))
 
     :componentWillReceiveProps
     (fn [props]
diff --git a/test/testcloact.cljs b/test/testcloact.cljs
index 78b4e1d..823f170 100644
--- a/test/testcloact.cljs
+++ b/test/testcloact.cljs
@@ -71,7 +71,7 @@
   (when isClient
     (let [ran (atom 0)
           comp (reagent/create-class
-                {:get-initial-state (fn [])
+                {:get-initial-state (fn [] {:foo "initial"})
                  :render
                  (fn []
                    (let [this (reagent/current-component)]
@@ -80,12 +80,14 @@
       (with-mounted-component (comp)
         (fn [C div]
           (swap! ran inc)
-          (is (found-in #"hi " div))
+          (is (found-in #"hi initial" div))
 
-          (reagent/set-state C {:foo "there"})
+          (reagent/replace-state C {:foo "there"})
+          (rflush)
           (is (found-in #"hi there" div))
 
           (reagent/set-state C {:foo "you"})
+          (rflush)
           (is (found-in #"hi you" div))))
       (is (= 4 @ran)))))