aoc ml

aoc.ml

type command =
  | Up of int
  | Down of int
  | Forward of int

type p1 = {
  horiz : int;
  depth : int;
}

let cmd1 { horiz ; depth } = function
  | Up n      -> { horiz             ; depth=(depth - n) }
  | Down n    -> { horiz             ; depth=(depth + n) }
  | Forward n -> { horiz=(horiz + n) ; depth };;

type p2 = {
  aim   : int;
  horiz : int;
  depth : int;
}

let cmd2 { horiz ; depth ; aim } = function
  | Up n      -> { horiz         ; depth               ; aim=aim-n }
  | Down n    -> { horiz         ; depth               ; aim=aim+n }
  | Forward n -> { horiz=horiz+n ; depth=depth+(aim*n) ; aim }

let cmd_constructors =
  Hashtbl.of_seq 
    (List.to_seq 
       [("up", fun x -> (Up x));
        ("down", fun x -> (Down x));
        ("forward", fun x -> (Forward x))])

let parse line = 
  Scanf.sscanf line "%s %i" 
    (fun s i -> (Hashtbl.find cmd_constructors s) i)

let fold_stream initial fn stream =
  let acc = ref initial in
  Stream.iter (fun line -> acc := (fn !acc line)) stream;
  !acc

let fold_p1 = fold_stream {horiz=0; depth=0} cmd1
let fold_p2 = fold_stream {horiz=0; depth=0; aim=0} cmd2 

let command_stream channel =
    Stream.from
      (fun _ ->
         try Some (parse (input_line channel))
         with End_of_file -> None)

let process_file filename fn =
  let in_channel = open_in filename in
  let res = ref None in
  try
    res := Some (fn (command_stream in_channel));
    close_in in_channel;
    !res
  with e ->
    close_in in_channel;
    raise e;

let s1 f = match (process_file f fold_p1) with
  | None -> None
  | Some {horiz; depth} -> Some (horiz * depth)

let s2 f = match (process_file f fold_p2) with
  | None -> None
  | Some {horiz; depth} -> Some (horiz * depth)

let res = [s1 "/tmp/a.test";
           s2 "/tmp/a.test"]