Solve day 4 part 1

day_4.sml

@@ -1,4 +1,4 @@
-fun amountInStr (needle : string, haystack : string) : int =
+fun amountInStr (needle, haystack) : int =
   let
     val l = List.tabulate (size haystack, (fn i =>
       String.isPrefix needle (String.substring (haystack, i, (size haystack - i)))
@@ -7,36 +7,65 @@ fun amountInStr (needle : string, haystack : string) : int =
     foldl (fn (v, acc) => if v then acc + 1 else acc) 0 l
   end
 
-fun reverseStr (str : string) = (String.implode o List.rev o String.explode) str
+fun reverseStr str = (String.implode o List.rev o String.explode) str
 
-fun wordsInLine (line : string) =
+fun wordsInLine line =
   amountInStr ("XMAS", line) + amountInStr ("XMAS", reverseStr line)
 
 fun wordsInLines (lines : string list) =
   foldl (op +) 0 (map wordsInLine lines)
 
-fun columns (lines : string list) =
-  let
-    fun lineLength (i : int) = size (List.nth (lines, i))
+fun lineLength lines i = size (List.nth (lines, i))
 
-    fun column (i : int) =
-      String.implode (List.tabulate (
-        lineLength i,
-        (fn j => String.sub (List.nth (lines, j), i))
-      ))
-  in
-    List.tabulate (length lines, column)
-  end
+fun pos lines (x, y) = String.sub (List.nth (lines, y), x)
+
+fun column lines i =
+  String.implode (List.tabulate (
+    lineLength lines i,
+    (fn j => pos lines (i, j))
+  ))
+
+fun nwDiagonal lines i =
+  String.implode (List.tabulate (i + 1, (fn j =>
+    pos lines (i - j, j)
+  )))
+
+fun seDiagonal lines i =
+  String.implode (List.tabulate (i, (fn j =>
+    pos lines ((lineLength lines i) - (i - j), (lineLength lines i) - j - 1)
+  )))
+
+fun neDiagonal lines i =
+  String.implode (List.tabulate (i, (fn j =>
+    pos lines ((lineLength lines i) - (i - j), j)
+  )))
+
+fun swDiagonal lines i =
+  String.implode (List.tabulate (i + 1, (fn j =>
+    pos lines (i - j, (lineLength lines i) - j - 1)
+  )))
+
+fun applyToLines f lines =
+  List.tabulate (length lines, (f lines))
 
 val file = TextIO.openIn "day_4.txt"
 val input = TextIO.inputAll file;
 TextIO.closeIn file;
 
-val lines = String.tokens (fn (char) => char = #"\n") input
+val lines = String.tokens (fn c => c = #"\n") input
 
-val columns = columns lines
+val columns = applyToLines column lines
+
+val nwDiagonals = applyToLines nwDiagonal lines
+val seDiagonals = applyToLines seDiagonal lines
+val neDiagonals = applyToLines neDiagonal lines
+val swDiagonals = applyToLines swDiagonal lines
 
 val result =
   wordsInLines lines +
-  wordsInLines columns
+  wordsInLines columns +
+  wordsInLines nwDiagonals +
+  wordsInLines seDiagonals +
+  wordsInLines neDiagonals +
+  wordsInLines swDiagonals