sml-projects/rope.sml

signature ROPE = sig
  type t
  val empty : t
  val of_string : string -> t
  val size : t -> int
  val insert : int * string * t -> t
end

structure Rope :> ROPE = struct
  datatype t  
    = N0 of string
    | N1 of t
    | N2 of t * int * t
    | L2 of string * string
    | N3 of t * t * t

  datatype balance  
    = AddedNode
    | DeletedNode
    | NoAction

  val target_length = 1024
  val empty = N0 ""
  fun of_string string = N0 string

  fun is_less_than_target(str1, str2) =
    String.size str1 + String.size str2 <= target_length

  exception AuxConstructor

  local
    fun help_size(acc, rope) =
      case rope of
        N0 s => 
          acc + String.size s
      | N1 t => 
          help_size(acc, t)
      | N2(_, lm, r) =>
          help_size(acc + lm, r)
      | _ => raise AuxConstructor
  in
    fun size rope = help_size(0, rope)
  end

  fun ins_root rope =
    case rope of
      L2(s1, s2) => 
        N2(N0 s1, String.size s1, N0 s2)
    | N3(t1, t2, t3) =>
        let
          val t1_size = size t1
          val left = N2(t1, t1_size, t2)
          val t2_size = size t2
        in
          N2(left, t1_size + t2_size, N1 t3)
        end
    | t =>
        N1 t

  fun del_root rope =
    case rope of
      N1 t => t
    | t => t

  fun ins_n1 rope =
    case rope of
        L2 (s1, s2) =>
          N2(N0 s1, String.size s1, N0 s2)
      | N3(t1, t2, t3) =>
          let 
            val left = N2(t1, size t1, t2)
          in
            N2(left, size left, t2)
          end
      | t =>
          N1 t

  fun ins_n2_left (left, right) =
    case (left, right) of
      (L2(s1, s2), t3) =>
        N3(N0 s1, N0 s2, t3)
    | (N3(t1, t2, t3), N1 t4) =>
        let
          val t1_size = size t1
          val left = N2(t1, t1_size, t2)
          val t3_size = size t3
          val right = N2(t3, t3_size, t4)
          val t2_size = size t2
          val left_size = t1_size + t2_size
        in
          N2(left, left_size, right)
        end
    | (N3(t1, t2, t3), t4) =>
        let 
          val left = N2(t1, size t2, t2)
        in
          N3(left, N1 t3, t4)
        end
    | (l, r) =>
        N2(l, size l, r)

  fun del_n2_left(left, right) = 
    case (left, right) of
      (N1 t1, N1 t2) =>
        N1(N2(t1, size t1, t2))
    | (N1 (N1 t1), N2(N1 t2, _, (t3 as N2 _))) =>
        let
          val left = N2(t1, size t1, t2)
          val inner = N2(left, size left, t3)
        in
          N1 (inner)
        end
    | (N1 (N1 t1), N2(N2(t2, _, t3), _, N1 t4)) =>
        let
          val left = N2(t1, size t1, t2)
          val right = N2(t3, size t3, t4)
          val inner = N2(left, size left, right)
        in
          N1 inner
        end
    | (N1 (t1 as N1 _), N2((t2 as N2 _), _, (t3 as N2 _))) =>
        let
          val left = N2(t1, size t1, t2)
          val right = N1 t3
        in
          N2(left, size left, right)
        end
    | (l, r) =>
        N2(l, size l, r)

  fun ins_n2_right (left, right) =
    case (left, right) of
      (t1, L2(s1, s2)) =>
        N3(t1, N0 s1, N0 s2)
    | (N1 t1, N3(t2, t3, t4)) =>
        let
          val left = N2(t1, size t1, t2)
          val t3_size = size t3
          val right = N2(t3, t3_size, t4)
        in
          N2(left, size left, right)
        end
    | (t1, N3(t2, t3, t4)) =>
        let
          val right = N2(t3, size t3, t4)
        in
          N3(t1, N1 t2, right)
        end
    | (l, r) =>
        N2(l, size l, r)

  fun del_n2_right(left, right) =
    case (left, right) of
      (N2(N1 t1, _, N2(t2, _, t3)), N1 t4) =>
        let
          val left = N2(t1, size t1, t2)
          val right = N2(t3, size t3, t4)
          val inner = N2(left, size left, right)
        in
          N1 inner
        end
    | (N2((t1 as N2 _), lm, N1 t2), N1 (N1 t3)) =>
        let
          val right = N2(t2, size t2, t3)
          val inner = N2(t1, lm, right)
        in
          N1 inner
        end
    | (N2( (t1 as N2 _), _, (t2 as N2 _)), N1 (N1 t3)) =>
        let
          val left_size = size t1
          val left = N1 t1
          val right = N2(t2, size t2, t3)
        in
          N2(left, left_size, right)
        end
    | (l, r) =>
        N2(l, size l, r)

  local
    fun ins_leaf(cur_index, new_str, rope, old_str) =
      if cur_index <= 0 then
        if is_less_than_target(old_str, new_str) then
          (N0(new_str ^ old_str), NoAction)
        else
          (L2(new_str, old_str), AddedNode)
      else if cur_index >= String.size old_str then
        if is_less_than_target(old_str, new_str) then
          (N0(old_str ^ new_str), NoAction)
        else
          (L2(old_str, new_str), AddedNode)
      else
        (* Need to split in middle of string. *)
        let
          val sub1 = String.substring(old_str, 0, cur_index)
          val sub2_len = String.size old_str - cur_index
          val sub2 = String.substring(old_str, cur_index, sub2_len)
        in
          if is_less_than_target(old_str, new_str) then
            (N0(sub1 ^ new_str ^ sub2), NoAction)
          else if cur_index + String.size new_str <= target_length then
            (L2(sub1 ^ new_str, sub2), AddedNode)
          else if (String.size old_str - cur_index) + String.size new_str <= target_length then
            (L2(sub1, new_str ^ sub2), AddedNode)
          else
            (N3(N0 sub1, N0 new_str, N0 sub2), AddedNode)
        end
    fun ins(cur_index, new_str, rope) =
      case rope of
        N2(l, lm, r) =>
          if cur_index < lm then
            let 
              val (l, action) = ins(cur_index, new_str, l)
            in
              (case action of
                NoAction => 
                  (case (l, r) of
                    (N0 s1, N0 s2) =>
                      if is_less_than_target(s1, s2) then
                        (N0 (s1 ^ s2), DeletedNode)
                      else
                        (N2(l, lm + String.size new_str, r), action)
                  | _ =>
                    (N2(l, lm + String.size new_str, r), action))
              | AddedNode =>
                  (ins_n2_left(l, r), action)
              | DeletedNode =>
                  (del_n2_left(l, r), action))
            end
          else
            let
              val (r, action) = ins(cur_index - lm, new_str, r)
            in
              (case action of
                NoAction =>
                  (case (l, r) of
                    (N0 s1, N0 s2) =>
                      if is_less_than_target(s1, s2) then
                        (N0 (s1 ^ s2), DeletedNode)
                      else
                        (N2(l, lm, r), action)
                    | _ =>
                        (N2(l, lm, r), action))
                | AddedNode =>
                    (ins_n2_right(l, r), action)
                | DeletedNode =>
                    (del_n2_right(l, r), action))
            end
      | N1 t =>
          let
            val (t, action) = ins(cur_index, new_str, t)
          in
            (case action of
              AddedNode =>
                (ins_n1 t, action)
            | _ =>
                (N1 t, action))
          end
      | N0 old_str => 
          ins_leaf(cur_index, new_str, rope, old_str)
      | _ =>
          raise AuxConstructor
   in
    fun insert (index, str, rope) =
      let
        val (rope, action) = ins(index, str, rope)
      in
        (case action of
          NoAction =>
            rope
        | AddedNode =>
            ins_root rope
        | DeletedNode =>
            del_root rope)
      end
  end
end