(* ps2-explain.sml

   Solution to Problem 2 of Problem Set 2. *)

(* val shortest : (int -> bool) -> str -> str * str

   If w is an str of zero's and one's, and there is a prefix x of w
   such that f(diff x), then shortest f w returns (x, y), where x is
   the shortest such prefix and y is such that x @ y = w. *)

fun shortest f (w : str) : str * str =
      let fun short(bs, n, nil)     =
                if f n then (bs, nil) else raise Fail "shouldn't happen"
            | short(bs, n, c :: cs) =
                if f n
                then (bs, c :: cs)
                else short(bs @ [c], n + diffSym c, cs)
      in short(nil, 0, w) end

(* val shortestPositive : str -> str * str

   If w is an str of zero's and one's, and there is a prefix x of w
   such that diff x >= 1, then shortestPositive w returns (x, y),
   where x is the shortest such prefix and y is such that x @ y = w. *)

val shortestPositive = shortest(fn n => n >= 1)

(* val shortestNegative : str -> str * str

   If w is an str of zero's and one's, and there is a prefix x of w
   such that diff x <= ~1, then shortestNegative w returns (x, y),
   where x is the shortest such prefix and y is such that x @ y = w. *)

val shortestNegative = shortest(fn n => n <= ~1)

(* val splitPositive : str -> str * str

   If w is an str of zero's and one's and diff w >= 1, then
   splitPositive w returns a pair (x, y), such that w = x @ [zero] @ y,
   diff x = 0 and diff y = diff w - 1. *)

fun splitPositive w : str * str =
      let val (u, y) = shortestPositive w  (* w = u @ y *)
          (* diff u >= 1 and shorter prefixes of w have non-positive diffs *)
          val x      = Str.allButLast u  (* diff x = 0 *)
          val b      = Str.last u  (* isZero b *)
          (* u = x @ [zero], w = x @ [zero] @ y, diff y = diff w - 1 *)
      in (x, y) end          

(* val explain : str -> expl

   If w is in Y, then explain w returns an explanation expl such
   that Str.equal(strExplained expl, w). *)

fun explain (w : str) =
      if null w
        then Rule1
      else if isZero(hd w)
        then let val t      = tl w  (* w = [zero] @ t, diff t = ~1 *)
                 val (u, v) = shortestNegative t  (* t = u @ v *)
                 (* diff u <= ~1 and strictly shorter prefixes of t
                    than u have non-negative diffs *)
                 val x      = Str.allButLast u  (* diff x in {0, 1} *)
                 val b      = Str.last u  (* isOne b *)
                 (* w = [zero] @ x @ [one] @ v *)
             in if diff x = 0
                then (* diff v = 1 *)
                     let val (y, z) = splitPositive v  (* v = y @ [zero] @ z *)
                         (* diff y = 0, diff z = 0,
                            w = ([zero] @ x @ [one] @ y @ [zero]) @ z *)
                     in Rule5(Rule3(explain x, explain y), explain z) end
                else (* diff x = 1, diff v = 0, x begins with 0 *)
                     let val y = tl x  (* diff y = 0 *)
                         (* w = ([zero] @ nil @ [zero] @ y @ [one]) @ v *)
                     in Rule5(Rule2(Rule1, explain y), explain v) end
             end 
      else (* isOne(hd w) *)
           let val t      = tl w  (* w = [one] @ t, diff t = 2 *)
               val (x, u) = splitPositive t  (* t = x @ [zero] @ u,
                                                diff x = 0, diff u = 1 *)
               val (y, z) = splitPositive u  (* u = y @ [zero] @ z,
                                                diff y = 0, diff z = 0 *)
               (* w = ([one] @ x @ [zero] @ y @ [zero]) @ z *)
           in Rule5(Rule4(explain x, explain y), explain z) end;