module Representation.ParticipantTree

import Std.Foundation
import Std.List
import Std.Natural

-- A bounded-fan-out hierarchy over logical roles 0..count-1, numbered in
-- preorder so that every subtree is one contiguous role range [role,
-- role + size). Role 0 is the root. A node's descendants (its range minus
-- itself) are split into at most `fanout` contiguous chunks whose sizes differ
-- by at most one, larger chunks first; each chunk is one child's subtree and
-- its first role is that child.
--
-- Why preorder ranges: a parent certifies its children's coverage with O(fan-
-- out) checks -- consecutive child ranges must tile (role, role + size) -- so
-- a checkpoint or collective rooted here never collects the complete
-- membership or every shard descriptor at one participant
-- (docs/multi-device-learning/22-fleet-scale.md, "Distributed checkpoint data,
-- bounded publication metadata"). Locating a role descends one path from the
-- root, so its cost is depth x fan-out, not the participant count.
--
-- These are logical roles. Which device holds a role, and whether a
-- particular placement keeps parent and child in one failure domain, are
-- separate decisions (Execution.ResourceGraph); this relation only fixes who
-- exchanges with whom.
family ParticipantTreeNode : Type 0
constructor ParticipantTreeNodeOf
-- the node's role: the first role of its subtree
field unrestricted participantTreeRole : Nat
-- roles in the subtree, the node included
field unrestricted participantTreeSize : Nat
-- the parent's role; the root names itself
field unrestricted participantTreeParent : Nat
-- edges from the root to this node
field unrestricted participantTreeDepth : Nat

end-family

def participantTreeRoleOf =
  (lambda unrestricted node : (family ParticipantTreeNode) .
    (eliminate
      ParticipantTreeNode
      (lambda unrestricted current : (family ParticipantTreeNode) . Nat)
      node
      (branch ParticipantTreeNodeOf role size parent depth . role)))

def participantTreeSizeOf =
  (lambda unrestricted node : (family ParticipantTreeNode) .
    (eliminate
      ParticipantTreeNode
      (lambda unrestricted current : (family ParticipantTreeNode) . Nat)
      node
      (branch ParticipantTreeNodeOf role size parent depth . size)))

def participantTreeParentOf =
  (lambda unrestricted node : (family ParticipantTreeNode) .
    (eliminate
      ParticipantTreeNode
      (lambda unrestricted current : (family ParticipantTreeNode) . Nat)
      node
      (branch ParticipantTreeNodeOf role size parent depth . parent)))

def participantTreeDepthOf =
  (lambda unrestricted node : (family ParticipantTreeNode) .
    (eliminate
      ParticipantTreeNode
      (lambda unrestricted current : (family ParticipantTreeNode) . Nat)
      node
      (branch ParticipantTreeNodeOf role size parent depth . depth)))

-- The children of a node with `size` roles: min(fanout, size - 1).
def participantTreeChildCount =
  (lambda unrestricted fanout : Nat .
    (lambda unrestricted size : Nat . (naturalMinimum fanout (naturalPredecessor size))))

-- The size of child `index` (from 0): the descendants divided as evenly as
-- possible, the remainder going to the first children.
def participantTreeChildSize =
  (lambda unrestricted fanout : Nat .
    (lambda unrestricted size : Nat .
      (lambda unrestricted index : Nat .
        (let unrestricted below =
          (naturalPredecessor size)
          in
          (let unrestricted children =
            (participantTreeChildCount fanout size)
            in
            (naturalAdd
              (naturalDivideUnchecked below children)
              (naturalSelect (naturalLess index (naturalModuloUnchecked below children)) 1 0)))))))

-- The first role of child `index`: the node, then every earlier child.
def participantTreeChildStart =
  (lambda unrestricted fanout : Nat .
    (lambda unrestricted role : Nat .
      (lambda unrestricted size : Nat .
        (lambda unrestricted index : Nat .
          (let unrestricted below =
            (naturalPredecessor size)
            in
            (let unrestricted children =
              (participantTreeChildCount fanout size)
              in
              (let unrestricted base =
                (naturalDivideUnchecked below children)
                in
                (naturalAdd
                  (succ role)
                  (naturalAdd
                    (naturalMultiply index base)
                    (naturalMinimum index (naturalModuloUnchecked below children)))))))))))

def participantTreeChildNode =
  (lambda unrestricted fanout : Nat .
    (lambda unrestricted node : (family ParticipantTreeNode) .
      (lambda unrestricted index : Nat .
        (constructor
          ParticipantTreeNode
          ParticipantTreeNodeOf
          (participantTreeChildStart
            fanout
            (participantTreeRoleOf node)
            (participantTreeSizeOf node)
            index)
          (participantTreeChildSize fanout (participantTreeSizeOf node) index)
          (participantTreeRoleOf node)
          (succ (participantTreeDepthOf node))))))

def participantTreeRoot =
  (lambda unrestricted count : Nat .
    (constructor ParticipantTreeNode ParticipantTreeNodeOf 0 count 0 0))

-- The index of the child whose range holds `target` (a descendant of node).
def participantTreeChildHolding =
  (lambda unrestricted fanout : Nat .
    (lambda unrestricted node : (family ParticipantTreeNode) .
      (lambda unrestricted target : Nat .
        (stdListFold
          Nat
          Nat
          (lambda unrestricted index : Nat .
            (lambda unrestricted rest : Nat .
              (naturalSelect
                (naturalLess
                  target
                  (participantTreeChildStart
                    fanout
                    (participantTreeRoleOf node)
                    (participantTreeSizeOf node)
                    (succ index)))
                index
                rest)))
          0
          (stdListNaturalsBelow (participantTreeChildCount fanout (participantTreeSizeOf node)))))))

-- Descend from `node` toward `target`, at most `fuel` edges. Fan-out of at
-- least two bounds the depth by log2 of the count, so 64 levels cover every
-- role a 64-bit count names; a chain (fan-out one) needs fuel = count.
def participantTreeDescend =
  (lambda unrestricted fanout : Nat .
    (lambda unrestricted fuel : Nat .
      (nat-eliminate
        (lambda unrestricted remaining : Nat .
          (pi unrestricted node : (family ParticipantTreeNode) .
            (pi unrestricted target : Nat . (family ParticipantTreeNode))))
        (lambda unrestricted node : (family ParticipantTreeNode) .
          (lambda unrestricted target : Nat . node))
        (lambda unrestricted previous : Nat .
          (lambda unrestricted deeper : (pi unrestricted node : (family ParticipantTreeNode) . (pi unrestricted target : Nat . (family ParticipantTreeNode))) .
            (lambda unrestricted node : (family ParticipantTreeNode) .
              (lambda unrestricted target : Nat .
                (nat-eliminate
                  (lambda unrestricted found : Nat . (family ParticipantTreeNode))
                  (deeper
                    (participantTreeChildNode
                      fanout
                      node
                      (participantTreeChildHolding fanout node target))
                    target)
                  (lambda unrestricted predecessor : Nat .
                    (lambda unrestricted ignored : (family ParticipantTreeNode) . node))
                  (naturalEqual (participantTreeRoleOf node) target))))))
        fuel)))

def participantTreeFuel =
  (lambda unrestricted fanout : Nat .
    (lambda unrestricted count : Nat . (naturalSelect (naturalLess fanout 2) count 64)))

-- The node of `role`. Admission (participantTreeAdmitted) must hold.
def participantTreeNode =
  (lambda unrestricted fanout : Nat .
    (lambda unrestricted count : Nat .
      (lambda unrestricted role : Nat .
        (participantTreeDescend
          fanout
          (participantTreeFuel fanout count)
          (participantTreeRoot count)
          role))))

-- The node's children, as nodes, in role order.
def participantTreeChildren =
  (lambda unrestricted fanout : Nat .
    (lambda unrestricted node : (family ParticipantTreeNode) .
      (stdListMap
        Nat
        (family ParticipantTreeNode)
        (participantTreeChildNode fanout node)
        (stdListNaturalsBelow (participantTreeChildCount fanout (participantTreeSizeOf node))))))

def participantTreeIsRoot =
  (lambda unrestricted node : (family ParticipantTreeNode) .
    (naturalIsZero (participantTreeDepthOf node)))

-- A hierarchy exists for a nonempty count with a nonzero fan-out, and the
-- role must be one of its participants.
def participantTreeAdmitted =
  (lambda unrestricted fanout : Nat .
    (lambda unrestricted count : Nat .
      (lambda unrestricted role : Nat .
        (naturalAnd
          (naturalNonzero fanout)
          (naturalAnd (naturalNonzero count) (naturalLess role count))))))

-- The deepest node: the largest subtree's chain of first children reaches it,
-- because sizes are nonincreasing in child order.
def participantTreeHeight =
  (lambda unrestricted fanout : Nat .
    (lambda unrestricted count : Nat .
      (participantTreeDepthOf
        (nat-eliminate
          (lambda unrestricted remaining : Nat . (family ParticipantTreeNode))
          (participantTreeRoot count)
          (lambda unrestricted previous : Nat .
            (lambda unrestricted node : (family ParticipantTreeNode) .
              (nat-eliminate
                (lambda unrestricted leaf : Nat . (family ParticipantTreeNode))
                node
                (lambda unrestricted predecessor : Nat .
                  (lambda unrestricted ignored : (family ParticipantTreeNode) .
                    (participantTreeChildNode fanout node 0)))
                (naturalLess 1 (participantTreeSizeOf node)))))
          (participantTreeFuel fanout count)))))

-- The nodes from the root down to (excluding) `role`'s node: the ancestors
-- a participant consults to verify that its own subtree is part of one
-- published root, in root-first order. At most the hierarchy's height.
def participantTreeAncestors =
  (lambda unrestricted fanout : Nat .
    (lambda unrestricted count : Nat .
      (lambda unrestricted role : Nat .
        (app
          (nat-eliminate
            (lambda unrestricted remaining : Nat .
              (pi unrestricted node : (family ParticipantTreeNode) .
                (family StdList (family ParticipantTreeNode))))
            (lambda unrestricted node : (family ParticipantTreeNode) .
              (constructor StdList StdListEmpty (family ParticipantTreeNode)))
            (lambda unrestricted previous : Nat .
              (lambda unrestricted deeper : (pi unrestricted node : (family ParticipantTreeNode) . (family StdList (family ParticipantTreeNode))) .
                (lambda unrestricted node : (family ParticipantTreeNode) .
                  (nat-eliminate
                    (lambda unrestricted found : Nat .
                      (family StdList (family ParticipantTreeNode)))
                    (constructor
                      StdList
                      StdListCons
                      (family ParticipantTreeNode)
                      node
                      (deeper
                        (participantTreeChildNode
                          fanout
                          node
                          (participantTreeChildHolding fanout node role))))
                    (lambda unrestricted predecessor : Nat .
                      (lambda unrestricted ignored : (family StdList (family ParticipantTreeNode)) .
                        (constructor StdList StdListEmpty (family ParticipantTreeNode))))
                    (naturalEqual (participantTreeRoleOf node) role)))))
            (participantTreeFuel fanout count))
          (participantTreeRoot count)))))

-- The index, among `node`'s children, of the child whose subtree holds `role`.
def participantTreeChildIndex =
  participantTreeChildHolding