path: root/stdlib/source/library/lux/data/collection/dictionary/ordered.lux

(.using
 [library
  [lux {"-" has revised}
   [abstract
    equivalence
    [monad {"+" Monad do}]
    ["[0]" order {"+" Order}]]
   [control
    ["[0]" maybe]]
   [data
    ["p" product]
    [collection
     ["[0]" list ("[1]#[0]" monoid mix)]]]
   [macro
    ["^" pattern]]
   [math
    [number
     ["n" nat]]]]])

(def: error_message
  "Invariant violation")

(type: Color
  (Variant
   {#Red}
   {#Black}))

(type: (Node k v)
  (Record
   [#color Color
    #key k
    #value v
    #left (Maybe (Node k v))
    #right (Maybe (Node k v))]))

(template [<create> <color>]
  [(def: (<create> key value left right)
     (All (_ k v) (-> k v (Maybe (Node k v)) (Maybe (Node k v)) (Node k v)))
     [#color {<color>}
      #key key
      #value value
      #left left
      #right right])]

  [red   #Red]
  [black #Black]
  )

(type: .public (Dictionary k v)
  (Record
   [#order (Order k)
    #root (Maybe (Node k v))]))

(def: .public (empty order)
  (All (_ k v) (-> (Order k) (Dictionary k v)))
  [#order order
   #root {.#None}])

... TODO: Doing inneficient access of Order functions due to compiler bug.
... TODO: Must improve it as soon as bug is fixed.
(def: .public (value key dict)
  (All (_ k v) (-> k (Dictionary k v) (Maybe v)))
  (let [... (open "_#[0]") (the #order dict)
        ]
    (loop (again [node (the #root dict)])
      (case node
        {.#None}
        {.#None}

        {.#Some node}
        (let [node_key (the #key node)]
          (cond (# dict = node_key key)
                ... (_#= node_key key)
                {.#Some (the #value node)}

                (# dict < node_key key)
                ... (_#< node_key key)
                (again (the #left node))

                ... (_#> (the #key node) key)
                (again (the #right node))))
        ))))

... TODO: Doing inneficient access of Order functions due to compiler bug.
... TODO: Must improve it as soon as bug is fixed.
(def: .public (key? dict key)
  (All (_ k v) (-> (Dictionary k v) k Bit))
  (let [... (open "_#[0]") (the #order dict)
        ]
    (loop (again [node (the #root dict)])
      (case node
        {.#None}
        #0

        {.#Some node}
        (let [node_key (the #key node)]
          (or (# dict = node_key key)
              ... (_#= node_key key)
              (if (# dict < node_key key)
                ... (_#< node_key key)
                (again (the #left node))
                (again (the #right node)))))))))

(template [<name> <side>]
  [(def: .public (<name> dict)
     (All (_ k v) (-> (Dictionary k v) (Maybe v)))
     (case (the #root dict)
       {.#None}
       {.#None}

       {.#Some node}
       (loop (again [node node])
         (case (the <side> node)
           {.#None}
           {.#Some (the #value node)}

           {.#Some side}
           (again side)))))]

  [min #left]
  [max #right]
  )

(def: .public (size dict)
  (All (_ k v) (-> (Dictionary k v) Nat))
  (loop (again [node (the #root dict)])
    (case node
      {.#None}
      0

      {.#Some node}
      (++ (n.+ (again (the #left node))
               (again (the #right node)))))))

(def: .public empty?
  (All (_ k v) (-> (Dictionary k v) Bit))
  (|>> ..size (n.= 0)))

(template [<name> <other_color> <self_color> <no_change>]
  [(def: (<name> self)
     (All (_ k v) (-> (Node k v) (Node k v)))
     (case (the #color self)
       {<other_color>}
       (.has #color {<self_color>} self)

       {<self_color>}
       <no_change>
       ))]

  [blackened #Red   #Black self]
  [reddened  #Black #Red   (panic! error_message)]
  )

(def: (with_left addition center)
  (All (_ k v) (-> (Node k v) (Node k v) (Node k v)))
  (case (the #color center)
    {#Red}
    (red (the #key center)
         (the #value center)
         {.#Some addition}
         (the #right center))
    
    {#Black}
    (with_expansions
      [<default_behavior> (these (black (the #key center)
                                        (the #value center)
                                        {.#Some addition}
                                        (the #right center)))]
      (case (the #color addition)
        {#Red}
        (case (the #left addition)
          (^.multi {.#Some left}
                   [(the #color left) {#Red}])
          (red (the #key addition)
               (the #value addition)
               {.#Some (blackened left)}
               {.#Some (black (the #key center)
                              (the #value center)
                              (the #right addition)
                              (the #right center))})

          _
          (case (the #right addition)
            (^.multi {.#Some right}
                     [(the #color right) {#Red}])
            (red (the #key right)
                 (the #value right)
                 {.#Some (black (the #key addition)
                                (the #value addition)
                                (the #left addition)
                                (the #left right))}
                 {.#Some (black (the #key center)
                                (the #value center)
                                (the #right right)
                                (the #right center))})

            _
            <default_behavior>))
        
        {#Black}
        <default_behavior>))))

(def: (with_right addition center)
  (All (_ k v) (-> (Node k v) (Node k v) (Node k v)))
  (case (the #color center)
    {#Red}
    (red (the #key center)
         (the #value center)
         (the #left center)
         {.#Some addition})
    
    {#Black}
    (with_expansions
      [<default_behavior> (these (black (the #key center)
                                        (the #value center)
                                        (the #left center)
                                        {.#Some addition}))]
      (case (the #color addition)
        {#Red}
        (case (the #right addition)
          (^.multi {.#Some right}
                   [(the #color right) {#Red}])
          (red (the #key addition)
               (the #value addition)
               {.#Some (black (the #key center)
                              (the #value center)
                              (the #left center)
                              (the #left addition))}
               {.#Some (blackened right)})

          _
          (case (the #left addition)
            (^.multi {.#Some left}
                     [(the #color left) {#Red}])
            (red (the #key left)
                 (the #value left)
                 {.#Some (black (the #key center)
                                (the #value center)
                                (the #left center)
                                (the #left left))}
                 {.#Some (black (the #key addition)
                                (the #value addition)
                                (the #right left)
                                (the #right addition))})

            _
            <default_behavior>))
        
        {#Black}
        <default_behavior>))))

(def: .public (has key value dict)
  (All (_ k v) (-> k v (Dictionary k v) (Dictionary k v)))
  (let [(open "_#[0]") (the #order dict)
        root' (loop (again [?root (the #root dict)])
                (case ?root
                  {.#None}
                  {.#Some (red key value {.#None} {.#None})}

                  {.#Some root}
                  (let [reference (the #key root)]
                    (`` (cond (~~ (template [<comp> <tag> <add>]
                                    [(<comp> reference key)
                                     (let [side_root (the <tag> root)
                                           outcome (again side_root)]
                                       (if (same? side_root outcome)
                                         ?root
                                         {.#Some (<add> (maybe.trusted outcome)
                                                        root)}))]

                                    [_#<                             #left  ..with_left]
                                    [(order.> (the #order dict)) #right ..with_right]
                                    ))

                              ... (_#= reference key)
                              {.#Some (.has #value value root)}
                              )))
                  ))]
    (.has #root root' dict)))

(def: (left_balanced key value ?left ?right)
  (All (_ k v) (-> k v (Maybe (Node k v)) (Maybe (Node k v)) (Node k v)))
  (case ?left
    (^.multi {.#Some left}
             [(the #color left) {#Red}]
             [(the #left left) {.#Some left>>left}]
             [(the #color left>>left) {#Red}])
    (red (the #key left)
         (the #value left)
         {.#Some (blackened left>>left)}
         {.#Some (black key value (the #right left) ?right)})

    (^.multi {.#Some left}
             [(the #color left) {#Red}]
             [(the #right left) {.#Some left>>right}]
             [(the #color left>>right) {#Red}])
    (red (the #key left>>right)
         (the #value left>>right)
         {.#Some (black (the #key left)
                        (the #value left)
                        (the #left left)
                        (the #left left>>right))}
         {.#Some (black key value
                        (the #right left>>right)
                        ?right)})

    _
    (black key value ?left ?right)))

(def: (right_balanced key value ?left ?right)
  (All (_ k v) (-> k v (Maybe (Node k v)) (Maybe (Node k v)) (Node k v)))
  (case ?right
    (^.multi {.#Some right}
             [(the #color right) {#Red}]
             [(the #right right) {.#Some right>>right}]
             [(the #color right>>right) {#Red}])
    (red (the #key right)
         (the #value right)
         {.#Some (black key value ?left (the #left right))}
         {.#Some (blackened right>>right)})

    (^.multi {.#Some right}
             [(the #color right) {#Red}]
             [(the #left right) {.#Some right>>left}]
             [(the #color right>>left) {#Red}])
    (red (the #key right>>left)
         (the #value right>>left)
         {.#Some (black key value ?left (the #left right>>left))}
         {.#Some (black (the #key right)
                        (the #value right)
                        (the #right right>>left)
                        (the #right right))})

    _
    (black key value ?left ?right)))

(def: (without_left key value ?left ?right)
  (All (_ k v) (-> k v (Maybe (Node k v)) (Maybe (Node k v)) (Node k v)))
  (case ?left
    (^.multi {.#Some left}
             [(the #color left) {#Red}])
    (red key value {.#Some (blackened left)} ?right)

    _
    (case ?right
      (^.multi {.#Some right}
               [(the #color right) {#Black}])
      (right_balanced key value ?left {.#Some (reddened right)})

      (^.multi {.#Some right}
               [(the #color right) {#Red}]
               [(the #left right) {.#Some right>>left}]
               [(the #color right>>left) {#Black}])
      (red (the #key right>>left)
           (the #value right>>left)
           {.#Some (black key value ?left (the #left right>>left))}
           {.#Some (right_balanced (the #key right)
                                   (the #value right)
                                   (the #right right>>left)
                                   (# maybe.functor each reddened (the #right right)))})

      _
      (panic! error_message))
    ))

(def: (without_right key value ?left ?right)
  (All (_ k v) (-> k v (Maybe (Node k v)) (Maybe (Node k v)) (Node k v)))
  (case ?right
    (^.multi {.#Some right}
             [(the #color right) {#Red}])
    (red key value ?left {.#Some (blackened right)})

    _
    (case ?left
      (^.multi {.#Some left}
               [(the #color left) {#Black}])
      (left_balanced key value {.#Some (reddened left)} ?right)

      (^.multi {.#Some left}
               [(the #color left) {#Red}]
               [(the #right left) {.#Some left>>right}]
               [(the #color left>>right) {#Black}])
      (red (the #key left>>right)
           (the #value left>>right)
           {.#Some (left_balanced (the #key left)
                                  (the #value left)
                                  (# maybe.functor each reddened (the #left left))
                                  (the #left left>>right))}
           {.#Some (black key value (the #right left>>right) ?right)})

      _
      (panic! error_message)
      )))

(def: (prepended ?left ?right)
  (All (_ k v) (-> (Maybe (Node k v)) (Maybe (Node k v)) (Maybe (Node k v))))
  (case [?left ?right]
    [{.#None} _]
    ?right

    [_ {.#None}]
    ?left

    [{.#Some left} {.#Some right}]
    (case [(the #color left) (the #color right)]
      [{#Red} {#Red}]
      (do maybe.monad
        [fused (prepended (the #right left) (the #right right))]
        (case (the #color fused)
          {#Red}
          (in (red (the #key fused)
                   (the #value fused)
                   {.#Some (red (the #key left)
                                (the #value left)
                                (the #left left)
                                (the #left fused))}
                   {.#Some (red (the #key right)
                                (the #value right)
                                (the #right fused)
                                (the #right right))}))

          {#Black}
          (in (red (the #key left)
                   (the #value left)
                   (the #left left)
                   {.#Some (red (the #key right)
                                (the #value right)
                                {.#Some fused}
                                (the #right right))}))))

      [{#Red} {#Black}]
      {.#Some (red (the #key left)
                   (the #value left)
                   (the #left left)
                   (prepended (the #right left)
                              ?right))}
      
      [{#Black} {#Red}]
      {.#Some (red (the #key right)
                   (the #value right)
                   (prepended ?left
                              (the #left right))
                   (the #right right))}

      [{#Black} {#Black}]
      (do maybe.monad
        [fused (prepended (the #right left) (the #left right))]
        (case (the #color fused)
          {#Red}
          (in (red (the #key fused)
                   (the #value fused)
                   {.#Some (black (the #key left)
                                  (the #value left)
                                  (the #left left)
                                  (the #left fused))}
                   {.#Some (black (the #key right)
                                  (the #value right)
                                  (the #right fused)
                                  (the #right right))}))
          
          {#Black}
          (in (without_left (the #key left)
                            (the #value left)
                            (the #left left)
                            {.#Some (black (the #key right)
                                           (the #value right)
                                           {.#Some fused}
                                           (the #right right))}))
          ))
      )

    _
    (undefined)))

(def: .public (lacks key dict)
  (All (_ k v) (-> k (Dictionary k v) (Dictionary k v)))
  (let [(open "_#[0]") (the #order dict)
        [?root found?] (loop (again [?root (the #root dict)])
                         (case ?root
                           {.#Some root}
                           (let [root_key (the #key root)
                                 root_val (the #value root)]
                             (if (_#= root_key key)
                               [(prepended (the #left root)
                                           (the #right root))
                                #1]
                               (let [go_left? (_#< root_key key)]
                                 (case (again (if go_left?
                                                (the #left root)
                                                (the #right root)))
                                   [{.#None} #0]
                                   [{.#None} #0]

                                   [side_outcome _]
                                   (if go_left?
                                     (case (the #left root)
                                       (^.multi {.#Some left}
                                                [(the #color left) {#Black}])
                                       [{.#Some (without_left root_key root_val side_outcome (the #right root))}
                                        #0]

                                       _
                                       [{.#Some (red root_key root_val side_outcome (the #right root))}
                                        #0])
                                     (case (the #right root)
                                       (^.multi {.#Some right}
                                                [(the #color right) {#Black}])
                                       [{.#Some (without_right root_key root_val (the #left root) side_outcome)}
                                        #0]

                                       _
                                       [{.#Some (red root_key root_val (the #left root) side_outcome)}
                                        #0])
                                     )))
                               ))

                           {.#None}
                           [{.#None} #0]
                           ))]
    (case ?root
      {.#None}
      (if found?
        (.has #root ?root dict)
        dict)

      {.#Some root}
      (.has #root {.#Some (blackened root)} dict)
      )))

(def: .public (revised key transform dict)
  (All (_ k v) (-> k (-> v v) (Dictionary k v) (Dictionary k v)))
  (case (..value key dict)
    {.#Some old}
    (..has key (transform old) dict)

    {.#None}
    dict))

(def: .public (of_list order list)
  (All (_ k v) (-> (Order k) (List [k v]) (Dictionary k v)))
  (list#mix (function (_ [key value] dict)
              (..has key value dict))
            (empty order)
            list))

(template [<name> <type> <output>]
  [(def: .public (<name> dict)
     (All (_ k v) (-> (Dictionary k v) (List <type>)))
     (loop (again [node (the #root dict)])
       (case node
         {.#None}
         (list)

         {.#Some node'}
         (all list#composite
              (again (the #left node'))
              (list <output>)
              (again (the #right node'))))))]

  [entries [k v] [(the #key node') (the #value node')]]
  [keys    k     (the #key node')]
  [values  v     (the #value node')]
  )

(implementation: .public (equivalence (open ",#[0]"))
  (All (_ k v) (-> (Equivalence v) (Equivalence (Dictionary k v))))
  
  (def: (= reference sample)
    (let [(open "/#[0]") (the #order reference)]
      (loop (again [entriesR (entries reference)
                    entriesS (entries sample)])
        (case [entriesR entriesS]
          [{.#End} {.#End}]
          #1

          [{.#Item [keyR valueR] entriesR'} {.#Item [keyS valueS] entriesS'}]
          (and (/#= keyR keyS)
               (,#= valueR valueS)
               (again entriesR' entriesS'))

          _
          #0)))))