:- pred constr(bool). 
:- mode constr(in). 
:- ignore constr/1.
:- type tree(D) ---> leaf ; node( D,   tree(D),     tree(D) ). 
%                 empty tree     key  left-subtree right-subtree
% tree(d) is the type of the trees whose keys have type D. 
% Here we assume that "D" is "int". Note: duplicates are NOT allowed.     
% ==========================================================
% Program. 
% https://algs4.cs.princeton.edu/32bst/

:- pred bstdelete(int,tree(int),tree(int)).
:- mode bstdelete(in,in,out).

bstdelete(X,leaf,leaf).                           % note deleting X from empty tree succeeds.
bstdelete(X,node(A,leaf,R), R) :- constr(X=A).
bstdelete(X,node(A,L,leaf), L) :- constr(X=A).
bstdelete(X,node(A,L,R), node(Succ,L,T) ) :-  constr(X=A & IsDef), 
    bstdeletemin(R, T, Succ, IsDef).  

bstdelete(X,node(A,L,R), node(A,L1,R)) :- constr(X<A), bstdelete(X,L,L1).
bstdelete(X,node(A,L,R), node(A,L,R1)) :- constr(A<X), bstdelete(X,R,R1).

% ---
:- pred bstdeletemin(tree(int),tree(int),int,bool).
:- mode bstdeletemin(in,out,out,out).

bstdeletemin(leaf,leaf,X,IsDef) :- constr(X=0 & ~IsDef).       % 0 can be any int value.
bstdeletemin(node(A,leaf,R),R,X,IsDef) :- constr(X=A & IsDef). % non-linear head is ok
bstdeletemin(node(A,node(LA,LL,LR),R),node(A,T,R),X,IsDef) :- constr(IsDef),
        bstdeletemin(node(LA,LL,LR),T,X,IsDef). 

% X is the successor of A in the tree 
% X is the minimum (and the leftmost node) of the right subtree rooted at A.

% ==========================================================
% catamorphisms
:- pred bstree(tree(int),bool).
:- mode bstree(in,out).
:- cata bstree/2-1.

bstree(leaf,B) :- constr(B).   % B = true
bstree(node(N,L,R),B) :- 
  bstree(L,BL), 
  bstree(R,BR),
  treemax(L,IsDefL,MaxL), 
  treemin(R,IsDefR,MinR),
  constr( ( (GeqLeft=(IsDefL => N>MaxL)) &
          (LeqRight=(IsDefR => N<MinR)) ) &
          (B = ((GeqLeft & BL) & (LeqRight & BR)))).

:- pred treemax(tree(int),bool,int).  
:- mode treemax(in,out,out).
:- cata treemax/3-1.

treemax(leaf,IsDef,Max) :- constr(~IsDef & Max=0).
treemax(node(N,L,R),IsDef,Max) :- 
  treemax(L,IsDefL,MaxL),
  treemax(R,IsDefR,MaxR),
  constr(IsDef & 
         (Max= ite(IsDefL & IsDefR, 
                     ite(MaxL>MaxR, 
                          ite(MaxL>N,MaxL,N), 
                          ite(MaxR>N,MaxR,N) ),            % max(MaxL,max(MaxR,N))
                     ite(IsDefL,
                          ite(MaxL>N,MaxL,N),              % max(MaxL,N)
                          ite(IsDefR & MaxR>N,MaxR,N))))). % max(MaxR,N)
% ----
:- pred treemin(tree(int),bool,int). 
:- mode treemin(in,out,out).
:- cata treemin/3-1.

treemin(leaf,IsDef,Min) :- constr(~IsDef & Min=0). 
treemin(node(N,L,R),IsDef,Min) :- 
    treemin(L,IsDefL,MinL),
    treemin(R,IsDefR,MinR),
  constr(IsDef & 
         (Min= ite(IsDefL & IsDefR, 
                     ite(MinL<MinR, 
                          ite(MinL<N,MinL,N), 
                          ite(MinR<N,MinR,N) ),            % min(MinL,min(MinR,N))
                     ite(IsDefL,
                          ite(MinL<N,MinL,N),              % min(MinL,N)
                          ite(IsDefR & MinR<N,MinR,N))))). % min(MinR,N)
%-----  COUNT --- how many X's in the given tree.
:- pred treecount(int,tree(int),int).
:- mode treecount(in,in,out).
:- cata treecount/3-2.

treecount(X,leaf,Res) :- constr( Res=0 ).
treecount(X,node(A,L,R),Res) :- 
    constr( Res=ite(X=A,ResL+ResR+1,ResL+ResR) ),
    treecount(X,L,ResL),
    treecount(X,R,ResR).   
    
% ==========================================================
% Contracts on bstdelete:
:- spec bstdelete(X,T1,T2) ==>
     bstree(T1,B1),
     treecount(Y,T1,C1), treecount(Y,T2,C2)
     =>
     constr( B1 => C2=ite(X=Y, 0, C1 ) ).

% the contract on bstdelete is NOT checked  
% without the following contract on bstdeletemin
:- spec bstdeletemin(T1,T2,Min,IsDef) ==>
     bstree(T1,B1), treecount(Y,T1,C1), treecount(Y,T2,C2)
     =>
     constr( (IsDef & B1) => C2=ite(Min=Y, 0, C1 ) ).
     
% ==========================================================
% Verification. 
:- pred ff1.
ff1 :- 
     constr(~ ( B1 => C2=ite(X=Y, 0, C1 ) )),
     bstree(T1,B1), treecount(Y,T1,C1), treecount(Y,T2,C2),    
     bstdelete(X,T1,T2).
     
:- pred ff2.
ff2 :-
    constr( ~( (IsDef & B1) => C2=ite(Min=Y, 0, C1 ) )),
    bstree(T1,B1), treecount(Y,T1,C1), treecount(Y,T2,C2),
    bstdeletemin(T1,T2,Min,IsDef).
% ==========================================================





% NON-CATA predicates: bstdelete bstdeletemin constr