justify/Main.lean
mehbark 0706170d95 remove Array.replicate and #evals
`Array.mkArray` does the same thing. couldn't remember the name!
2025-01-28 11:31:50 -05:00

72 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
-- 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 String.replicate : Nat → Char → String
| 0, _ => ""
| n+1, x => replicate n x |>.push x
def sep (xs space : Array String)
: String := Id.run do
let mut out := ""
let mut i := 0
while h : i < xs.size do
out := out ++ xs[i]
if let some space := space[i]?
then out := out ++ space
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 smallSpaces := numSpaces - largeSpaces
let largeSpace := smallSpace.push ' '
let spaces := Array.mkArray largeSpaces largeSpace
++ Array.mkArray smallSpaces smallSpace
sep words spaces
def breakUp (width : Nat) (words : Array String) : Array (Array String) :=
let (_, line, lines) := words.foldl
(λ(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] =>
if let some width := width.toNat?
then do
let mut text := ""
let stdin ← IO.getStdin
while true do
let line ← stdin.getLine
if line.isEmpty then break
text := text ++ line
IO.println $ solve width text
return 0
else fail
| _ => fail