From 5057b7772bc3d4e9a9461bdd49406c41ba86c2a8 Mon Sep 17 00:00:00 2001
From: Aaron Smith <83140718+centau@users.noreply.github.com>
Date: Thu, 3 Aug 2023 12:04:47 +0100
Subject: [PATCH]

---
 src/spring.luau | 19 +++++++++++++++++--
 test/tests.luau | 25 +++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/src/spring.luau b/src/spring.luau
index 2c44b0c..d33ea5a 100644
--- a/src/spring.luau
+++ b/src/spring.luau
@@ -36,10 +36,12 @@ type SpringData<T> = {
     duration: number,
     period: number,
     damping_ratio: number,
+
     velocity: number,
     initial_velocity: number,
     initial_position: T,
     target_position: T,
+
     target_updated: boolean,
     target: () -> T
 }
@@ -110,9 +112,10 @@ local function spring<T>(target: () -> T, period: number?, damping_ratio: number
         target = target
     }
 
-    local function input_changed()
+    local function input_changed(node)
         data.target_updated = true
         data.target_position = target()
+        springs[data] = node
     end
 
     for _, input in next, inputs do
@@ -124,11 +127,13 @@ local function spring<T>(target: () -> T, period: number?, damping_ratio: number
     return output_get
 end
 
+local remove_queue = {}
+
 local function update_springs(dt: number)
     for data, output in next, springs do
         if data.target_updated then
             data.target_updated = false
-            data.target = data.target()
+            data.target_position = data.target()
             data.initial_position = get(output)
             data.alpha = 0
             data.duration = 0
@@ -168,7 +173,17 @@ local function update_springs(dt: number)
         local value = lerp(initial_position, target_position, new_alpha)
 
         set(output, value)
+
+        if new_alpha > 0.9 then
+            table.insert(remove_queue, data)
+        end
     end
+
+    for _, data in next, remove_queue do
+        springs[data] = nil
+    end
+
+    table.clear(remove_queue)
 end
 
 return function() return spring, update_springs end
diff --git a/test/tests.luau b/test/tests.luau
index b4536ee..b8cd10e 100644
--- a/test/tests.luau
+++ b/test/tests.luau
@@ -1016,6 +1016,7 @@ TEST("spring()", function()
     local create = vide.create
     local source = vide.source
     local spring = vide.spring
+    local watch = vide.watch
 
     do CASE "Update state (on next hearbeat resumption cycle)"
         local value = source(10)
@@ -1067,6 +1068,30 @@ TEST("spring()", function()
         gc()
         CHECK(wref[1]) -- `output` should not gc
     end
+
+    do CASE "Spring finished"
+        local input = source(0)
+        local output = spring(input)
+
+        input(1)
+        vide.step(1e9) -- spring alpha at ~1
+
+        local count = -1
+        watch(function()
+            output()
+            count += 1
+        end)
+
+        vide.step(1) -- spring should be internally removed from spring queue
+        CHECK(count == 0)
+
+        --
+
+        gc() -- perform full gc
+        input(2) -- spring should be re-added to spring queue
+        vide.step(0) -- process spring queue
+        CHECK(count == 1) -- check spring was rescheduled correctly
+    end
 end)
 
 TEST("Events", function()