diff --git a/Main.lean b/Main.lean
index 71f1ed7..cc36042 100644
--- a/Main.lean
+++ b/Main.lean
@@ -3,12 +3,8 @@ def String.words (s : String) :=
 def String.lines (s : String) :=
   s.splitOn "\n" |>.filter (¬·.isEmpty) |>.toArray
 
--- TODO: remove trailing
-def Array.unlines (xs : Array String) : String := Id.run do
-  let mut out := ""
-  for line in xs do
-    out := out ++ line ++ "\n"
-  return out
+def Array.unlines (xs : Array String) : String :=
+  xs.foldl (· ++ · ++ "\n") ""
 
 def String.replicate : Nat → Char → String
 | 0,   _ => ""
@@ -25,22 +21,21 @@ def sep (xs : Array String) (space : Fin (xs.size-1) → String)
     i := i + 1
   return out
 
-def justify : Nat → Array String → String
-| _,   #[word] => word
-| width, words =>
-    let wordsLength   := words.foldr (·.length + ·) 0
-    let spacesNeeded  := width - wordsLength
-    let numSpaces     := words.size - 1
-    let smallSpaceLen := spacesNeeded / numSpaces
-    let smallSpace    := String.replicate smallSpaceLen ' '
-    let largeSpaces   := spacesNeeded % numSpaces
-    let largeSpace    := smallSpace.push ' '
-    let space i := if i.val ≤ largeSpaces then largeSpace else smallSpace
-    sep words space
+def justify (width : Nat) (words : Array String) : String :=
+  sep words space
+  where
+  space i := if i.val ≤ largeSpaces then largeSpace else smallSpace
+  largeSpace    := smallSpace.push ' '
+  smallSpace    := String.replicate smallSpaceLen ' '
+  smallSpaceLen := spacesNeeded / numSpaces
+  largeSpaces   := spacesNeeded % numSpaces
+  numSpaces     := words.size - 1
+  spacesNeeded  := width - wordsLength
+  wordsLength   := words.foldr (·.length + ·) 0
 
 def breakUp (width : Nat) (words : Array String) : Array (Array String) :=
   let (_, line, lines) := words.foldl
-    (λ(left, line, lines) word ↦
+    (fun (left, line, lines) word =>
       if word.length < left
       then (left - (word.length + 1), line.push word, lines)
       else (width - word.length, #[word], lines.push line))
@@ -54,20 +49,20 @@ def usage := "Usage: justify WIDTH"
 def fail : IO UInt32 := do IO.println usage; return 1
 
 def main : List String → IO UInt32
-| [width] =>
+| [width] => do
   if let some width := width.toNat?
-  then do
+  then
     let stdin ← IO.getStdin
     let mut remainder : Array String := #[]
-    while true do
+    repeat
       let line ← stdin.getLine
       if line.isEmpty then break
       let lines := breakUp width (remainder ++ line.words)
       if h : lines.size > 0 then
         remainder := lines[lines.size-1]
         for line in lines[0:lines.size-1] do
-          IO.println $ justify width line
-    IO.println $ String.intercalate " " remainder.toList
+          IO.println <| justify width line
+    IO.println <| sep remainder (fun _ => " ")
     return 0
   else fail
 | _ => fail
diff --git a/lakefile.toml b/lakefile.toml
index 59fa958..4d0bd25 100644
--- a/lakefile.toml
+++ b/lakefile.toml
@@ -5,3 +5,8 @@ defaultTargets = ["justify"]
 [[lean_exe]]
 name = "justify"
 root = "Main"
+
+[[require]]
+name = "batteries"
+scope = "leanprover-community"
+rev = "main"
diff --git a/lean-toolchain b/lean-toolchain
index f1f4005..3ca992c 100644
--- a/lean-toolchain
+++ b/lean-toolchain
@@ -1 +1 @@
-leanprover/lean4:4.15.0
+leanprover/lean4:v4.17.0-rc1