Sequence type in Unison supporting most operations in log_32(n) or O(1)

sequence.u

use data Array

structural type Sequence a
  = Empty Nat
  | Sequence Nat Nat (Array a) (Sequence (Array a)) (Array a)

Sequence.arity : Sequence a -> Nat
Sequence.arity = cases
  Sequence.Empty arity   -> arity
  Sequence arity _ _ _ _ -> arity

test> Sequence.tests = test.random do
  Each.repeat 100
  N = Random.natIn 0 100
  arity = Random.natIn 2 65
  nsl = List.range 0 N
  ns = Sequence.fromList arity nsl
  ns2 = List.foldRight Sequence.cons (Sequence.empty arity) nsl
  ns3 = 
    (l, r) = halve nsl 
    sl = List.foldRight Sequence.cons (Sequence.empty arity) l
    List.foldLeft (s a -> Sequence.snoc a s) sl r
  ensure (Sequence.toList ns === nsl) 
  ensure (Sequence.toList ns === Sequence.toList ns2)
  ensure (Sequence.toList ns3 === nsl) 
  ensure (List.map (i -> Sequence.unsafeAt i ns) nsl === nsl)
  ensure (List.map (i -> Sequence.unsafeAt i ns2) nsl === nsl)
  ensure (List.map (i -> Sequence.unsafeAt i ns3) nsl === nsl)
  ensure (List.unfold ns Sequence.uncons === nsl)
  ensure let
    step s = match Sequence.unsnoc s with
      None -> None
      Some (init, last) -> Some (last, init) 
    List.unfold ns step === List.reverse nsl
  k = Each.each nsl

  equalSeq s l = 
    ensureEqual (Sequence.toList s) l
    ensureEqual (List.map (i -> Sequence.unsafeAt i s) (List.range 0 (Sequence.size s))) l

  equalSeq (Sequence.append (Sequence.take k ns) (Sequence.drop k ns)) nsl
  equalSeq (Sequence.append (Sequence.take k ns) (Sequence.drop k ns2)) nsl
  equalSeq (Sequence.append (Sequence.take k ns) (Sequence.drop k ns3)) nsl
  equalSeq (Sequence.take k ns) (List.take k nsl)
  equalSeq (Sequence.take k ns2) (List.take k nsl)
  equalSeq (Sequence.take k ns3) (List.take k nsl)
  equalSeq (Sequence.drop k ns) (List.drop k nsl)
  equalSeq (Sequence.drop k ns2) (List.drop k nsl)
  equalSeq (Sequence.drop k ns3) (List.drop k nsl)

Sequence.append : Sequence a -> Sequence a -> Sequence a
Sequence.append s1 s2 = 
  if Sequence.size s1 >= Sequence.size s2 
  then Sequence.foldLeft (s a -> Sequence.snoc a s) s1 s2
  else Sequence.foldRight Sequence.cons s2 s1

Sequence.unsafeAt : Nat -> Sequence a -> a
Sequence.unsafeAt i = cases
  Sequence arity n hds mid tls 
    | i Nat.< Array.size hds            -> Array.unsafeAt i hds
    | i Nat.>= (n Nat.- Array.size tls) ->
      use Nat -
      Array.unsafeAt (i - (n - Array.size tls)) tls
    | otherwise                        ->
      use Nat - /
      i' = i - Array.size hds
      child = i' / arity
      j = Nat.mod i' arity
      Array.unsafeAt j (Sequence.unsafeAt child mid)

Sequence.cons : a -> Sequence a ->{} Sequence a
Sequence.cons a = cases
  Sequence.Empty arity ->
    Sequence arity 1 (Array.singleton a) (Sequence.Empty arity) Array.empty 
  Sequence arity n hds mid tls ->
    use Nat + ==
    if Array.size hds == arity then
      Sequence arity (n + 1) (Array.singleton a) (Sequence.cons hds mid) tls
    else
      Sequence arity (n + 1) (Array.cons a hds) mid tls

Sequence.empty : Nat -> Sequence a
Sequence.empty arity =
  use Nat <
  if arity < 2 then bug ("arity must be >= 2", 2) 
  else Sequence.Empty arity

Sequence.first : Sequence a ->{} Optional a
Sequence.first = cases
  Sequence.Empty _        -> None
  Sequence _ sz ls mid rs ->
    match Array.at 0 ls with
      Some a -> Some a
      _     ->
        use List head
        match Sequence.first mid with
          None    -> Array.at 0 rs
          Some as -> Array.at 0 as

Sequence.foldLeft : (b -> a ->{g} b) -> b -> Sequence a ->{g} b
Sequence.foldLeft f z = cases
  Sequence.Empty _ -> z
  Sequence _ _ hds mid tls ->
    z1 = Array.foldLeft f z hds
    z2 = Sequence.foldLeft (z as -> Array.foldLeft f z as) z1 mid
    z3 = Array.foldLeft f z2 tls
    z3

Sequence.foldRight : (a -> b ->{g} b) -> b -> Sequence a ->{g} b
Sequence.foldRight f z = cases
  Sequence.Empty _ -> z
  Sequence _ _ hds mid tls ->
    z1 = Array.foldRight f z tls
    z2 = Sequence.foldRight (as z -> Array.foldRight f z as) z1 mid
    z3 = Array.foldRight f z2 hds
    z3

Sequence.fromList : Nat -> [a] ->{} Sequence a
Sequence.fromList arity as = snocs as (Sequence.empty arity)

Sequence.last : forall a . Sequence a ->{} Optional a
Sequence.last = cases
  Sequence.Empty _        -> None
  Sequence _ sz ls mid rs ->
    match Array.at (Array.size rs - 1) rs with
      Some a -> Some a
      _     ->
        match Sequence.last mid with
          None    -> Array.at (Array.size ls - 1) ls
          Some as -> Array.at (Array.size as - 1) as

Sequence.size : Sequence a -> Nat
Sequence.size = cases
  Sequence.Empty _   -> 0
  Sequence _ n _ _ _ -> n

Sequence.snoc : a -> Sequence a ->{} Sequence a
Sequence.snoc a = cases
  Sequence.Empty arity ->
    Sequence arity 1 Array.empty (Sequence.Empty arity) (Array.singleton a)
  Sequence arity n hds mid tls ->
    use Nat + ==
    if Array.size tls == arity then
      Sequence arity (n + 1) hds (Sequence.snoc tls mid) (Array.singleton a)
    else Sequence arity (n + 1) hds mid (Array.snoc tls a)

Sequence.snocs : forall a . [a] -> Sequence a ->{} Sequence a
Sequence.snocs as b =
  n' =
    use Nat +
    Sequence.size b + Nat.min (List.size as) (Sequence.arity b)
  match as with
    [] -> b
    _ ->
      use Sequence Empty snocs
      use Array drop take
      match b with
        Empty arity ->
          b' = Sequence arity n' Array.empty (Empty arity) (Array.fromList (List.take arity as))
          snocs (List.drop arity as) b'
        Sequence arity n hds mid rs | Array.isEmpty rs ->
          b' = Sequence arity n' hds mid (Array.fromList (List.take arity as))
          snocs (List.drop arity as) b'
        _ -> List.foldLeft (b a -> Sequence.snoc a b) b as

Sequence.toList : forall a . Sequence a -> [a]
Sequence.toList = Sequence.foldLeft (:+) []

Sequence.uncons : forall a . Sequence a ->{} Optional (a, Sequence a)
Sequence.uncons = cases
  Sequence.Empty _ -> None
  Sequence arity sz hds mids tls -> match Array.size hds with
    0 -> match uncons mids with
      None 
        | Array.size tls == 0 -> None
        | otherwise ->    
          Some (Array.unsafeAt 0 tls, Sequence arity (sz - 1) (Array.dropClamped 1 tls) (Empty arity) Array.empty)  
      Some (hds, mid) -> 
        Some (Array.unsafeAt 0 hds, Sequence arity (sz - 1) (Array.dropClamped 1 hds) mid tls)
    n -> 
      Some (Array.unsafeAt 0 hds, Sequence arity (sz - 1) (Array.dropClamped 1 hds) mids tls)

Sequence.unsnoc : Sequence a ->{} Optional (Sequence a, a)
Sequence.unsnoc = cases
  Sequence.Empty _ -> None
  Sequence arity sz hds mids tls -> match Array.size tls with
    0 -> match unsnoc mids with
      None 
        | Array.size hds == 0 -> None
        | otherwise ->   
          Some (Sequence arity (sz - 1) Array.empty (Empty arity) (Array.dropRightClamped 1 hds), 
                Array.unsafeAt (Array.size hds - 1) hds)
      Some (mid, tls) -> 
        Some (Sequence arity (sz - 1) hds mid (Array.dropRightClamped 1 tls), Array.unsafeAt (arity - 1) tls)
    n -> 
      Some (Sequence arity (sz - 1) hds mids (Array.dropRightClamped 1 tls), Array.unsafeAt (Array.size tls - 1) tls)

Sequence.take : Nat -> Sequence a -> Sequence a
Sequence.take n s = match s with
  Sequence.Empty _ -> s 
  Sequence arity sz hds mids tls -> match n with
    0 -> Sequence.Empty arity
    _ | n >= sz -> s
    _ | n <= Array.size hds -> Sequence arity n (Array.take n hds) (Sequence.Empty arity) Array.empty 
    _ -> 
      hdMidSz = Array.size hds + (Sequence.size mids * arity)
      match n - hdMidSz with
        0 -> 
          k = n - Array.size hds
          mids' = Sequence.take (k+arity-1 / arity) mids
          if Nat.mod k arity == 0 then 
            Sequence arity n hds mids' Array.empty 
          else
            match Sequence.unsnoc mids' with
              Some (mids', last) -> Sequence arity n hds mids' (Array.take (k - (Sequence.size mids' * arity)) last) 
              None -> Sequence arity n hds mids' Array.empty
        k -> Sequence arity n hds mids (Array.take k tls)

-- drop the first n elements from `s`
-- implementation is analogous to take
Sequence.drop : Nat -> Sequence a -> Sequence a
Sequence.drop n s = match s with
  Sequence.Empty _ -> s 
  Sequence arity sz hds mids tls -> match n with
    0 -> s
    _ | n >= sz -> Sequence.Empty arity
    _ | n <= Array.size hds -> Sequence arity (sz - n) (Array.dropClamped n hds) mids tls 
    _ -> 
      hdMidSz = Array.size hds + (Sequence.size mids * arity)
      match n - hdMidSz with
        0 -> 
          k = n - Array.size hds
          mids' = Sequence.drop (k / arity) mids
          rem = Nat.mod k arity
          if rem == 0 then
            Sequence arity (sz - n) Array.empty mids' tls 
          else
            match Sequence.uncons mids' with
              Some (hd, mids') -> Sequence arity (sz - n) (Array.drop rem hd) mids' tls 
              None -> Sequence arity (sz - n) Array.empty mids' tls
        k -> Sequence arity (sz - n) (Array.drop k tls) (Empty arity) Array.empty

up.data.Array.takeClamped : Nat -> data.Array a -> data.Array a
up.data.Array.takeClamped n = cases
  Arr off len arr -> Arr off (Nat.min len n) arr

up.data.Array.dropClamped : Nat -> data.Array a -> data.Array a
up.data.Array.dropClamped n = cases
  Arr off len arr -> Arr (off + n) (len - n) arr

up.data.Array.dropRightClamped : Nat -> data.Array a -> data.Array a
up.data.Array.dropRightClamped n = cases
  Arr off len arr -> Arr off (len - n) arr