Fix allocations introduced in step (#454)

Author: devmotion

src/integrator.jl

@@ -226,10 +226,8 @@ function step(
     ϵ = fwd ? step_size(lf) : -step_size(lf)
     ϵ = ϵ'
 
-    res = if FullTraj
-        Vector{P}(undef, n_steps)
-    else
-        Vector{P}(undef, 1)
+    if FullTraj
+        res = Vector{P}(undef, n_steps)
     end
 
     (; θ, r) = z
@@ -252,20 +250,18 @@ function step(
         # Update result
         if FullTraj
             res[i] = z
-        else
-            res[1] = z
         end
         if !isfinite(z)
             # Remove undef
             if FullTraj
-                res = res[isassigned.(Ref(res), 1:n_steps)]
+                resize!(res, i)
             end
             break
         end
     end
     return if FullTraj
         res
     else
-        first(res)
+        z
     end
 end