fix dropJ function

1 files changed, 10 insertions, 7 deletions

diff --git a/lec07/JoinList.hs b/lec07/JoinList.hs
index 4eba3d2..25667c7 100644
--- a/lec07/JoinList.hs
+++ b/lec07/JoinList.hs
@@ -32,11 +32,14 @@ indexJ n (Append m jl1 jl2)
 
 dropJ :: (Sized b, Monoid b) => Int -> JoinList b a -> JoinList b a
 dropJ _ Empty = Empty
-dropJ n jl
-      | n <= 0 = jl
-dropJ n (Single _ _)
+dropJ n jl@(Single _ _)
       | n >= 1 = Empty
-dropJ n (Append m jl1 jl2)
-      | n >= (getSize $ size m) = Empty
-      | n >= s1 = jl2
-      where s1 = getSize $ size $ tag jl1
+      | otherwise = jl
+dropJ n jl@(Append m jl1 jl2)
+      | n <= 0 = jl
+      | n >= m' = Empty
+      | n < s1 = Append (Size $ m' - n) (dropJ n jl1) jl2
+      | n >= s1 = dropJ (n - s1) jl2
+      where
+        s1 = getSize $ size $ tag jl1
+        m' = getSize $ size m