justify/Main.lean
2025-02-04 10:10:36 -05:00

69 lines
2.1 KiB
Text

def String.words (s : String) :=
s.split Char.isWhitespace |>.filter (¬·.isEmpty) |>.toArray
def String.lines (s : String) :=
s.splitOn "\n" |>.filter (¬·.isEmpty) |>.toArray
def Array.unlines (xs : Array String) : String :=
xs.foldl (· ++ · ++ "\n") ""
def String.replicate : Nat → Char → String
| 0, _ => ""
| n+1, x => replicate n x |>.push x
def sep (xs : Array String) (space : Fin (xs.size-1) → String)
: String := Id.run do
let mut out := ""
let mut i := 0
while h : i < xs.size do
out := out ++ xs[i]
if h : i < xs.size - 1 then
out := out ++ space ⟨i, h⟩
i := i + 1
return out
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
(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))
(width, #[], #[])
lines.push line
def solve (width : Nat) :=
Array.unlines ∘ Array.map (justify width) ∘ breakUp width ∘ String.words
def usage := "Usage: justify WIDTH"
def fail : IO UInt32 := do IO.println usage; return 1
def main : List String → IO UInt32
| [width] => do
if let some width := width.toNat?
then
let stdin ← IO.getStdin
let mut remainder : Array String := #[]
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 <| sep remainder (fun _ => " ")
return 0
else fail
| _ => fail