A weird tree monad proposed on reddit

strange-monad.agda

-- a proof of monad laws for a strange monad defined in:
-- https://www.reddit.com/r/haskell/comments/cb1j40/is_there_a_valid_monad_instance_for_binary_trees/etcqsts/

data TREE (a : Set) : Set where
  Leaf : a → TREE a
  Tree : a → TREE a → TREE a → TREE a

distribute : ∀ {a} → (TREE a) -> (TREE a) -> (TREE a) -> TREE a
distribute l r (Leaf a) = Tree a l r
distribute l r (Tree a cl cr) = Tree a (distribute l r cl) (distribute l r cr)

return : ∀ {a} → a → TREE a
return = Leaf

_>>=_ : ∀ {a b} → TREE a → (a → TREE b) → TREE b
Leaf a >>= f = f a
Tree a l r >>= f = distribute  (l >>= f) (r >>= f) (f a)

data _≡_ {A : Set} (x : A) : A → Set where
  refl : x ≡ x

infix 3 _≡_

trans : ∀ {A} {x y z : A} → x ≡ y → y ≡ z → x ≡ z
trans refl refl = refl

cong : ∀ {A B} (f : A → B) → ∀ {x y} → x ≡ y → f x ≡ f y
cong f refl = refl

cong₂ : ∀ {A B C} (f : A → B → C) → ∀ {a b c d} → a ≡ b → c ≡ d → f a c ≡ f b d
cong₂ f refl refl = refl

lemma-case0 : ∀ {a b} l r x (g : a → TREE b)
  → distribute (l >>= g) (r >>= g) (Leaf x >>= g) ≡ distribute l r (Leaf x) >>= g
lemma-case0 = λ l r x g → refl


distribute-assoc : ∀ {a} (l r cl cr x : TREE a) → 
  distribute l r (distribute cl cr x)
  ≡
  distribute (distribute l r cl) (distribute l r cr) x
distribute-assoc l r cl cr (Leaf x) = refl
distribute-assoc l r cl cr (Tree x ccl ccr) =
  cong₂ (Tree x) (distribute-assoc l r cl cr ccl) (distribute-assoc l r cl cr ccr)

lemma :
  ∀ {a b} l r x (g : a → TREE b)
  → distribute (l >>= g) (r >>= g) (x >>= g) ≡ distribute l r x >>= g
lemma l r (Leaf x) g = refl
lemma l r (Tree x cl cr) g =
   trans (cong (distribute (l >>= g) (r >>= g)) (lemma-case0 cl cr x g))
     (trans (distribute-assoc _ _ _ _ _)
       (cong₂ (λ q w → distribute q w (g x))
       (lemma l r cl g)
       (lemma l r cr g)))


assoc : ∀ {a b c} (f : a → TREE b) (g : b → TREE c) x → x >>= (λ x → f x >>= g) ≡ (x >>= f) >>= g
assoc f g (Leaf x) = refl
assoc f g (Tree x l r) =
   trans ( cong₂ (λ q w → distribute q w (f x >>= g)) (assoc f g l) (assoc f g r))
   (lemma (l >>= f) (r >>= f) (f x) g)

unit-l : ∀ {a b} (f : a → TREE b) (x : a) → return x >>= f ≡ f x
unit-l f x = refl

unit-r : ∀ {a} (x : TREE a) → x >>= return ≡ x
unit-r (Leaf x) = refl
unit-r (Tree x l r) = cong₂ (Tree x) (unit-r _) (unit-r _)