90 lines
2.8 KiB
Text
90 lines
2.8 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 ex := "The heat bloomed\tin December\nas the carnival season\n kicked into gear.\nNearly helpless with sun and glare, I avoided Rio's brilliant\nsidewalks\n and glittering beaches,\npanting in dark corners\nand waiting out the inverted southern summer."
|
||
|
||
#eval ex.words
|
||
#eval ex.lines
|
||
#eval IO.println ex.lines.unlines
|
||
|
||
def Array.replicate : Nat → α → Array α
|
||
| 0, _ => #[]
|
||
| n+1, x => replicate n x |>.push x
|
||
|
||
theorem Array.replicate_size_eq_n : (Array.replicate n x).size = n := by
|
||
induction n <;> simp [Array.replicate, *]
|
||
|
||
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.append xs[i]
|
||
if let some space := space[i]?
|
||
then out := out.append 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.replicate largeSpaces largeSpace
|
||
++ Array.replicate smallSpaces smallSpace
|
||
sep words spaces
|
||
|
||
#eval justify 1 "Hello, world, how is it going?".words
|
||
#eval justify 1 "Hello, world, how is it going?".words |>.length
|
||
|
||
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
|
||
|
||
#eval IO.println $ solve 40 ex
|
||
|
||
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
|