Fix hot patching hook types (#4381)

* fix hot patching hook types

* fix clippy

* fix doc check

* just expose the subsecond jumptable

* prefer panicking, forcing more cold branches

* make get_jump_table unsafe

---------

Co-authored-by: Jonathan Kelley <jkelleyrtp@gmail.com>

packages/core/src/scope_context.rs

@@ -451,45 +451,66 @@ impl Scope {
     /// ```
     pub fn use_hook<State: Clone + 'static>(&self, initializer: impl FnOnce() -> State) -> State {
         let cur_hook = self.hook_index.get();
-        let mut hooks = self.hooks.try_borrow_mut().expect("The hook list is already borrowed: This error is likely caused by trying to use a hook inside a hook which violates the rules of hooks.");
 
-        if cur_hook >= hooks.len() {
-            Runtime::with(|rt| {
-                rt.while_not_rendering(|| {
-                    hooks.push(Box::new(initializer()));
-                });
-            })
-            .unwrap()
+        // The hook list works by keeping track of the current hook index and pushing the index forward
+        // while retrieving the hook value.
+        self.hook_index.set(cur_hook + 1);
+
+        let mut hooks = self.hooks
+            .try_borrow_mut()
+            .expect("The hook list is already borrowed: This error is likely caused by trying to use  hook inside a hook which violates the rules of hooks.");
+
+        // Try and retrieve the hook value if it exists
+        if let Some(existing) = self.use_hook_inner::<State>(&mut hooks, cur_hook) {
+            return existing;
         }
 
-        self.use_hook_inner::<State>(hooks, cur_hook)
+        // Otherwise, initialize the hook value. In debug mode, we allow hook types to change after a hot patch
+        self.push_hook_value(&mut hooks, cur_hook, initializer())
     }
 
-    // The interior version that gets monoorphized by the `State` type but not the `initializer` type.
+    // The interior version that gets monomorphized by the `State` type but not the `initializer` type.
     // This helps trim down binary sizes
     fn use_hook_inner<State: Clone + 'static>(
         &self,
-        hooks: std::cell::RefMut<Vec<Box<dyn std::any::Any>>>,
+        hooks: &mut Vec<Box<dyn std::any::Any>>,
         cur_hook: usize,
+    ) -> Option<State> {
+        hooks.get(cur_hook).and_then(|inn| {
+            let raw_ref: &dyn Any = inn.as_ref();
+            raw_ref.downcast_ref::<State>().cloned()
+        })
+    }
+
+    /// Push a new hook value or insert the value into the existing slot, warning if this is not after a hot patch
+    fn push_hook_value<State: Clone + 'static>(
+        &self,
+        hooks: &mut Vec<Box<dyn std::any::Any>>,
+        cur_hook: usize,
+        value: State,
     ) -> State {
-        hooks
-            .get(cur_hook)
-            .and_then(|inn| {
-                self.hook_index.set(cur_hook + 1);
-                let raw_ref: &dyn Any = inn.as_ref();
-                raw_ref.downcast_ref::<State>().cloned()
-            })
-            .expect(
-                r#"
-                Unable to retrieve the hook that was initialized at this index.
-                Consult the `rules of hooks` to understand how to use hooks properly.
+        // If this is a new hook, push it
+        if cur_hook >= hooks.len() {
+            hooks.push(Box::new(value.clone()));
+            return value;
+        }
 
-                You likely used the hook in a conditional. Hooks rely on consistent ordering between renders.
-                Functions prefixed with "use" should never be called conditionally.
+        // If we're in dev mode, we allow swapping hook values if the hook was initialized at this index
+        if cfg!(debug_assertions) && unsafe { subsecond::get_jump_table().is_some() } {
+            hooks[cur_hook] = Box::new(value.clone());
+            return value;
+        }
 
-                Help: Run `dx check` to look for check for some common hook errors.
-                "#,
-            )
+        // Otherwise, panic
+        panic!(
+            r#"Unable to retrieve the hook that was initialized at this index.
+                    Consult the `rules of hooks` to understand how to use hooks properly.
+
+                    You likely used the hook in a conditional. Hooks rely on consistent ordering between renders.
+                    Functions prefixed with "use" should never be called conditionally.
+
+                    Help: Run `dx check` to look for check for some common hook errors."#
+        );
     }
 
     pub fn push_before_render(&self, f: impl FnMut() + 'static) {

packages/subsecond/subsecond/src/lib.rs

@@ -278,10 +278,26 @@ pub fn call<O>(mut f: impl FnMut() -> O) -> O {
 // multithreading, it might become an issue.
 static APP_JUMP_TABLE: AtomicPtr<JumpTable> = AtomicPtr::new(std::ptr::null_mut());
 static HOTRELOAD_HANDLERS: Mutex<Vec<Arc<dyn Fn() + Send + Sync>>> = Mutex::new(Vec::new());
+
+/// Register a function that will be called whenever a patch is applied.
+///
+/// This handler will be run immediately after the patch library is loaded into the process and the
+/// JumpTable has been set.
 pub fn register_handler(handler: Arc<dyn Fn() + Send + Sync + 'static>) {
     HOTRELOAD_HANDLERS.lock().unwrap().push(handler);
 }
-fn get_jump_table() -> Option<&'static JumpTable> {
+
+/// Get the current jump table, if it exists.
+///
+/// This will return `None` if no jump table has been set yet.
+///
+/// # Safety
+///
+/// The `JumpTable` returned here is a pointer into a leaked box. While technically this reference is
+/// valid, we might change the implementation to invalidate the pointer between hotpatches.
+///
+/// You should only use this lifetime in temporary contexts - not *across* hotpatches!
+pub unsafe fn get_jump_table() -> Option<&'static JumpTable> {
     let ptr = APP_JUMP_TABLE.load(std::sync::atomic::Ordering::Relaxed);
     if ptr.is_null() {
         return None;
@@ -378,7 +394,7 @@ impl<A, M, F: HotFunction<A, M>> HotFn<A, M, F> {
         }
 
         let known_fn_ptr = <F as HotFunction<A, M>>::call_it as *const () as usize;
-        if let Some(jump_table) = get_jump_table() {
+        if let Some(jump_table) = unsafe { get_jump_table() } {
             if let Some(ptr) = jump_table.map.get(&(known_fn_ptr as u64)).cloned() {
                 return HotFnPtr(ptr);
             }