Solve day 3 part 2

day_3.sml

@@ -1,30 +1,63 @@
+datatype instruction = NOOP | DO | DONT | MUL of (int * int)
+
 fun numberAt (i : int) = fn (s : string) =>
   if i >= String.size s then NONE
   else SOME (Substring.takel Char.isDigit (Substring.extract (s, i, NONE)))
 
-fun substrToInt (n : substring) = valOf (Int.fromString (Substring.string n))
+fun startsWith (needle : string, haystack : string, pos : int) =
+  size haystack > pos + size needle andalso String.substring (haystack, pos, size needle) = needle
 
-fun sum (l : int list) : int = foldl (op +) 0 l
+fun substrToInt (n : substring) = valOf (Int.fromString (Substring.string n))
 
 val file = TextIO.openIn "day_3.txt"
 val input = TextIO.inputAll file;
 TextIO.closeIn file;
 
 val instructions = List.tabulate ((size input), (fn i =>
+  if startsWith ("do()", input, i) then DO
+  else if startsWith ("don't()", input, i) then DONT
+  else
     let
       val n1 = numberAt (i + 4) input
       val n2 = Option.mapPartial (fn n => numberAt (i + 5 + Substring.size n) input) n1
     in
-    if isSome n1 andalso isSome n2 andalso
-      String.substring (input, i, 4) = "mul(" andalso
-      String.sub (input, i + 4 + Substring.size (valOf n1)) = #"," andalso
-      String.sub (input, i + 5 + Substring.size (valOf n1) + Substring.size (valOf n2)) = #")"
-    then SOME (n1, n2) else NONE
-end))
+      case (n1, n2) of
+        (SOME(a), SOME(b)) =>
+          if
+            startsWith ("mul(", input, i) andalso
+            String.sub (input, i + 4 + Substring.size a) = #"," andalso
+            String.sub (input, i + 5 + Substring.size a + Substring.size b) = #")"
+          then MUL (substrToInt a, substrToInt b) else NOOP
+         | _ => NOOP
+    end
+))
 
-val instructions = List.mapPartial
-  (fn (num) => case num of SOME (SOME x, SOME y) => SOME (substrToInt x, substrToInt y) | _ => NONE)
-  instructions
+val instructions = List.filter (fn (ins) => ins <> NOOP) instructions
 
-val result = sum (map (op *) instructions)
+fun calculate (instructions : instruction vector) =
+  Vector.foldl (fn (ins, v) => case ins of MUL (a, b) => v + a * b | _ => v) 0 instructions
+
+(* Part 1 *)
+
+val mulOnlyResult = calculate (Vector.fromList instructions)
+
+(* Part 2 *)
+
+val filteredInstructions = Vector.mapi
+  (fn (i, ins) =>
+    let
+      val prevInstructions = List.rev (List.take (instructions, i))
+      val lastToggle = List.find (fn (ins) => ins = DO orelse ins = DONT) prevInstructions
+      val shouldMultiply = case lastToggle of
+        SOME (ins) => ins = DO
+          | NONE => true
+    in
+      case ins of
+        MUL (a, b) => if shouldMultiply then ins else NOOP
+         | _ => ins
+    end
+  )
+  (Vector.fromList instructions)
+
+val allInsResult = calculate filteredInstructions