--
--  Copyright (C) 2022-2025, Free Software Foundation, Inc.
--
--  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
--

pragma Ada_2022;

with SPARK.Containers.Parameter_Checks;
with SPARK.Big_Integers; use SPARK.Big_Integers;

private with SPARK.Containers.Functional.Maps;

generic
   type Element_Type is private;

   with
     function Equivalent_Elements
       (Left : Element_Type; Right : Element_Type) return Boolean is "=";

   Enable_Handling_Of_Equivalence : Boolean := True;
   --  This constant should only be set to False when no particular handling
   --  of equivalence over keys is needed, that is, Equivalent_Keys defines a
   --  key uniquely.

   --  Ghost lemmas used to prove that Equivalent_Elements is an equivalence
   --  relation.

   with
     procedure Equivalent_Elements_Reflexive (X : Element_Type) is null
     with Ghost => Static;
   with
     procedure Equivalent_Elements_Symmetric (X, Y : Element_Type) is null
     with Ghost => Static;
   with
     procedure Equivalent_Elements_Transitive (X, Y, Z : Element_Type) is null
     with Ghost => Static;

package SPARK.Containers.Functional.Multisets with
    Ghost      => SPARKlib_Logic,
    SPARK_Mode => On,
    Always_Terminates
is
   pragma
     Annotate
       (GNATcheck,
        Exempt_On,
        "Restrictions:No_Specification_Of_Aspect => Iterable",
        "The following usage of aspect Iterable has been reviewed"
        & "for compliance with GNATprove assumption"
        & " [SPARK_ITERABLE]");
   type Multiset is private
   with
     Default_Initial_Condition => (SPARKlib_Full => Is_Empty (Multiset)),
     Iterable                  =>
       (First       => Iter_First,
        Next        => Iter_Next,
        Has_Element => Iter_Has_Element,
        Element     => Iter_Element),
     Aggregate                 =>
       (Empty => Empty_Multiset, Add_Named => Aggr_Include),
     Annotate                  =>
       (GNATprove, Container_Aggregates, "Predefined_Maps");
   pragma
     Annotate
       (GNATcheck,
        Exempt_Off,
        "Restrictions:No_Specification_Of_Aspect => Iterable");
   --  Multisets are empty when default initialized.
   --  "For in" quantification over multisets should not be used.
   --  "For of" quantification over multisets iterates over elements.
   --  Occurrences in a multiset are counted modulo equivalence, the whole
   --  equivalence class is included/excluded at once. As equivalence classes
   --  might be infinite, quantification over elements of a finite multiset
   --  could be infinite. Thus, quantified expressions cannot be executed and
   --  should only be used in disabled ghost code. This is enforced by using
   --  the SPARKlib_Full assertion level.

   ----------------------
   -- Basic Operations --
   ----------------------

   function Nb_Occurence
     (Container : Multiset; Element : Element_Type) return Big_Natural
   with
     --  Get the number of occurence of Element in the Container

     Global   => null,
     Annotate => (GNATprove, Container_Aggregates, "Get");

   procedure Lemma_Nb_Occurence_Equivalent
     (Container : Multiset; Element_1, Element_2 : Element_Type)
     --  Nb_Occurence returns the same result on all equivalent elements
   with
     Ghost    => SPARKlib_Full,
     Global   => null,
     Annotate => (GNATprove, Automatic_Instantiation),
     Pre      =>
       Enable_Handling_Of_Equivalence
       and then Equivalent_Elements (Element_1, Element_2),
     Post     =>
       Nb_Occurence (Container, Element_1)
       = Nb_Occurence (Container, Element_2);

   function Contains
     (Container : Multiset; Element : Element_Type) return Boolean
   with
     --  Returns True iff Element occurs at least one in Container

     Global   => null,
     Post     =>
       (SPARKlib_Full =>
          Contains'Result = (Nb_Occurence (Container, Element) > 0)),
     Annotate => (GNATprove, Inline_For_Proof),
     Annotate => (GNATprove, Iterable_For_Proof, "Contains");

   function Choose (Container : Multiset) return Element_Type
   with
     --  Returns an element of the Multiset

     Global => null,
     Pre    => (SPARKlib_Defensive => not Is_Empty (Container)),
     Post   => (SPARKlib_Full => Contains (Container, Choose'Result));

   function Cardinality (Container : Multiset) return Big_Natural
   with
     --  The Cardinality of a Multiset is the number of equivalence classes in
     --  the Multiset taking into account the number of occurences of these
     --  equivalent classes.

     Global => null,
     Post   =>
       (SPARKlib_Full =>
          (for all Element of Container =>
             Nb_Occurence (Container, Element) <= Cardinality'Result));

   ------------------------
   -- Property Functions --
   ------------------------

   function "<=" (Left, Right : Multiset) return Boolean
   with
     --  Left <= Right if all the elements occur less in Left than in Right

     Global => null,
     Post   =>
       (SPARKlib_Full =>
          "<="'Result
          = (for all Item of Left =>
               Nb_Occurence (Left, Item) <= Nb_Occurence (Right, Item))
          and then
            (if "<="'Result then Cardinality (Left) <= Cardinality (Right)));

   function "=" (Left, Right : Multiset) return Boolean
   with
     Global => null,
     Post   =>
       (SPARKlib_Full =>
          "="'Result
          = ((for all Element of Left =>
                Nb_Occurence (Left, Element) = Nb_Occurence (Right, Element))
             and then
               (for all Element of Right =>
                  Nb_Occurence (Left, Element)
                  = Nb_Occurence (Right, Element)))
          and then
            (if "="'Result then (Cardinality (Left) = Cardinality (Right))));
   --  Two Multiset are equal if and only if all the Element_Type have the same
   --  number of occurences for both Multisets (possibly 0).

   pragma Warnings (Off, "unused variable ""Element""");
   function Is_Empty (Container : Multiset) return Boolean
   with
     Global => null,
     Post   =>
       (SPARKlib_Full =>
          Is_Empty'Result = (for all Element of Container => False)
          and then Is_Empty'Result = (Cardinality (Container) = 0));
   pragma Warnings (On, "unused variable ""Element""");
   --  A Multiset is empty if it has no element. In other words, if the number
   --  of occurence of all the elements is 0.

   function Equal_Except
     (Left : Multiset; Right : Multiset; Element : Element_Type) return Boolean
   with
     --  Check if all the elements of Left except the equivalence class of
     --  Element have the same number of occurences in Left and in Right and
     --  conversely.

     Global => null,
     Post   =>
       (SPARKlib_Full =>
          Equal_Except'Result
          = ((for all E of Left =>
                (if not Equivalent_Elements (E, Element)
                 then Nb_Occurence (Left, E) = Nb_Occurence (Right, E)))
             and then
               (for all E of Right =>
                  (if not Equivalent_Elements (E, Element)
                   then Nb_Occurence (Left, E) = Nb_Occurence (Right, E)))));

   ----------------------------
   -- Construction Functions --
   ----------------------------

   --  For better efficiency of both proofs and execution, avoid using
   --  construction functions in annotations and rather use property functions.

   function Empty_Multiset return Multiset
   with
     --  Returns an empty Multiset

     Global => null,
     Post   => (SPARKlib_Full => Is_Empty (Empty_Multiset'Result));

   function Add (Container : Multiset; Element : Element_Type) return Multiset
   with
     Global => null,
     Post   =>
       (SPARKlib_Full =>
          Contains (Add'Result, Element)
          and then Cardinality (Add'Result) = Cardinality (Container) + 1
          and then
            Nb_Occurence (Add'Result, Element)
            = Nb_Occurence (Container, Element) + 1
          and then Equal_Except (Container, Add'Result, Element));
   --  Returns Container with the number of occurences of the equivalence class
   --  of Element incremented by one.

   function Add
     (Container : Multiset; Element : Element_Type; Count : Big_Positive)
      return Multiset
   with
     Global => null,
     Post   =>
       (SPARKlib_Full =>
          Contains (Add'Result, Element)
          and then Cardinality (Add'Result) = Cardinality (Container) + Count
          and then
            Nb_Occurence (Add'Result, Element)
            = Nb_Occurence (Container, Element) + Count
          and then Equal_Except (Container, Add'Result, Element));
   --  Returns Container with the number of occurences of the equivalence class
   --  of Element incremented by Count.

   function Remove_All
     (Container : Multiset; Element : Element_Type) return Multiset
   with
     Global => null,
     Pre    => (SPARKlib_Defensive => Contains (Container, Element)),
     Post   =>
       (SPARKlib_Full =>
          not Contains (Remove_All'Result, Element)
          and then
            Cardinality (Remove_All'Result)
            = Cardinality (Container) - Nb_Occurence (Container, Element)
          and then Equal_Except (Container, Remove_All'Result, Element));
   --  Returns Container with no occurences of the equivalence class of Element

   function Remove
     (Container : Multiset; Element : Element_Type; Count : Big_Positive := 1)
      return Multiset
   with
     Global => null,
     Pre    =>
       (SPARKlib_Defensive => Count <= Nb_Occurence (Container, Element)),
     Post   =>
       (SPARKlib_Full =>
          Nb_Occurence (Remove'Result, Element)
          = Nb_Occurence (Container, Element) - Count
          and then
            Cardinality (Remove'Result) = Cardinality (Container) - Count
          and then Equal_Except (Container, Remove'Result, Element));
   --  Returns Container with the number of occurences of the equivalence class
   --  of Element decremented by Count.

   function Sum (Left : Multiset; Right : Multiset) return Multiset
   with
     Global => null,
     Post   =>
       (SPARKlib_Full =>
          Cardinality (Sum'Result) = Cardinality (Left) + Cardinality (Right)
          and Left <= Sum'Result
          and Right <= Sum'Result
          and
            (for all Element of Sum'Result =>
               Nb_Occurence (Sum'Result, Element)
               = Nb_Occurence (Left, Element)
                 + Nb_Occurence (Right, Element)));
   --  Returns the sum of Left and Right, in which the number of occurences of
   --  an element E is the sum of its number of occurences in Left and its
   --  number of occurences in Right.

   function "+" (Left : Multiset; Right : Multiset) return Multiset
   renames Sum;

   function Difference (Left : Multiset; Right : Multiset) return Multiset
   with
     --  Returns the difference of Left and Right, in which the number of
     --  occurences of an element E is the difference between its number of
     --  occurences in Left and its number of occurences in Right.

     Global => null,
     Post   =>
       (SPARKlib_Full =>
          (for all E of Left =>
             (if Nb_Occurence (Left, E) > Nb_Occurence (Right, E)
              then Contains (Difference'Result, E)))
          and then
            (for all E of Difference'Result =>
               Nb_Occurence (Difference'Result, E)
               = Max (0, Nb_Occurence (Left, E) - Nb_Occurence (Right, E))));

   function "-" (Left : Multiset; Right : Multiset) return Multiset
   renames Difference;

   function Intersection (Left : Multiset; Right : Multiset) return Multiset
   with
     --  Returns the intersection of Left and Right

     Global => null,
     Post   =>
       (SPARKlib_Full =>
          (for all Element of Left =>
             (if Contains (Right, Element)
              then Contains (Intersection'Result, Element)))
          and then
            (for all Element of Intersection'Result =>
               Nb_Occurence (Intersection'Result, Element)
               = Min
                   (Nb_Occurence (Left, Element),
                    Nb_Occurence (Right, Element))));

   function Union (Left : Multiset; Right : Multiset) return Multiset
   with
     --  Returns the union of Left and Right, i.e. the smallest Multiset
     --  containing both Left and Right.

     Global => null,
     Post   =>
       (SPARKlib_Full =>
          Left <= Union'Result
          and then Right <= Union'Result
          and then
            (for all Element of Union'Result =>
               Nb_Occurence (Union'Result, Element)
               = Max
                   (Nb_Occurence (Left, Element),
                    Nb_Occurence (Right, Element))));

   ------------
   -- Lemmas --
   ------------

   pragma Warnings (Off, "actuals for this call may be in wrong order");
   procedure Lemma_Sym_Intersection (Left : Multiset; Right : Multiset)
   with
     --  State that the Intersection is symmetrical

     Ghost  => SPARKlib_Full,
     Global => null,
     Post   => Intersection (Left, Right) = Intersection (Right, Left);

   procedure Lemma_Sym_Union (Left : Multiset; Right : Multiset)
   with
     --  State that the Union is symmetrical

     Ghost  => SPARKlib_Full,
     Global => null,
     Post   => Union (Left, Right) = Union (Right, Left);

   procedure Lemma_Sym_Sum (Left : Multiset; Right : Multiset)
   with
     --  State that the Sum is symmetrical.

     Ghost  => SPARKlib_Full,
     Global => null,
     Post   => Sum (Left, Right) = Sum (Right, Left);
   pragma Warnings (On, "actuals for this call may be in wrong order");

   ---------------------------------------
   -- Iteration on Functional Multisets --
   ---------------------------------------

   --  The Iterable aspect can be used to quantify over a functional multiset.
   --  However, if it is used to create a for loop, it will not allow users to
   --  prove their loops as there is no way to speak about the elements which
   --  have or have not been traversed already in a loop invariant. The
   --  function Iterate returns an object of a type Iterable_Multiset which can
   --  be used for iteration. The cursor is a functional multiset containing
   --  all the elements which have not been traversed yet. The current element
   --  being traversed being the result of Choose on this set.

   pragma
     Annotate
       (GNATcheck,
        Exempt_On,
        "Restrictions:No_Specification_Of_Aspect => Iterable",
        "The following usage of aspect Iterable has been reviewed"
        & "for compliance with GNATprove assumption"
        & " [SPARK_ITERABLE]");
   type Iterable_Multiset is private
   with
     Iterable =>
       (First       => First,
        Has_Element => Has_Element,
        Next        => Next,
        Element     => Element);
   pragma
     Annotate
       (GNATcheck,
        Exempt_Off,
        "Restrictions:No_Specification_Of_Aspect => Iterable");

   function Multiset_Logic_Equal (Left, Right : Multiset) return Boolean
   with Ghost => SPARKlib_Full, Annotate => (GNATprove, Logical_Equal);
   --  Logical equality on multisets

   function Iterate (Container : Multiset) return Iterable_Multiset
   with
     Global => null,
     Post   =>
       (SPARKlib_Full =>
          Multiset_Logic_Equal (Get_Multiset (Iterate'Result), Container));
   --  Return an iterator over a functional multiset

   function Get_Multiset (Iterator : Iterable_Multiset) return Multiset
   with Global => null;
   --  Retrieve the multiset associated with an iterator

   function Valid_Subset
     (Iterator : Iterable_Multiset; Cursor : Multiset) return Boolean
   with
     Global => null,
     Post   =>
       (SPARKlib_Full =>
          (if Valid_Subset'Result then Cursor <= Get_Multiset (Iterator)));
   --  Returns True on all multisets which can be reached by iterating over
   --  Container.

   function Element
     (Iterator : Iterable_Multiset; Cursor : Multiset) return Element_Type
   with
     Global   => null,
     Pre      => (SPARKlib_Defensive => not Is_Empty (Cursor)),
     Post     => (SPARKlib_Full => Element'Result = Choose (Cursor)),
     Annotate => (GNATprove, Inline_For_Proof);
   --  The next element to be considered for the iteration is the result of
   --  choose on Cursor.

   function First (Iterator : Iterable_Multiset) return Multiset
   with
     Global => null,
     Post   =>
       (SPARKlib_Full =>
          Multiset_Logic_Equal (First'Result, Get_Multiset (Iterator))
          and then Valid_Subset (Iterator, First'Result));
   --  In the first iteration, the cursor is the multiset associated with
   --  Iterator.

   function Next
     (Iterator : Iterable_Multiset; Cursor : Multiset) return Multiset
   with
     Global => null,
     Pre    =>
       (SPARKlib_Defensive =>
          Valid_Subset (Iterator, Cursor) and then not Is_Empty (Cursor)),
     Post   =>
       (SPARKlib_Full =>
          Valid_Subset (Iterator, Next'Result)
          and then
            Multiset_Logic_Equal
              (Next'Result, Remove_All (Cursor, Choose (Cursor))));
   --  At each iteration, remove the equivalence class of considered element
   --  from the Cursor set.

   function Has_Element
     (Iterator : Iterable_Multiset; Cursor : Multiset) return Boolean
   with
     Global => null,
     Post   =>
       (SPARKlib_Full =>
          Has_Element'Result
          = (Valid_Subset (Iterator, Cursor) and then not Is_Empty (Cursor)));
   --  Return True on non-empty sets which can be reached by iterating over
   --  Container.

   --------------------------
   -- Instantiation Checks --
   --------------------------

   --  Check that the actual parameters follow the appropriate assumptions

   function Copy_Element (Item : Element_Type) return Element_Type
   is (Item);
   --  Elements of containers are copied by numerous primitives in this
   --  package. This function causes GNATprove to verify that such a copy is
   --  valid (in particular, it does not break the ownership policy of SPARK,
   --  i.e. it does not contain pointers that could be used to alias mutable
   --  data).

   package Eq_Elements_Checks is new
     SPARK.Containers.Parameter_Checks.Equivalence_Checks
       (T                   => Element_Type,
        Eq                  => Equivalent_Elements,
        Param_Eq_Reflexive  => Equivalent_Elements_Reflexive,
        Param_Eq_Symmetric  => Equivalent_Elements_Symmetric,
        Param_Eq_Transitive => Equivalent_Elements_Transitive);
   --  Check that the actual parameter for Equivalent_Elements is an
   --  equivalence relation.

   --------------------------------------------------
   -- Iteration Primitives Used For Quantification --
   --------------------------------------------------

   type Private_Key is private;

   function Iter_First (Container : Multiset) return Private_Key
   with Ghost => SPARKlib_Full, Global => null;

   function Iter_Has_Element
     (Container : Multiset; Key : Private_Key) return Boolean
   with Ghost => SPARKlib_Full, Global => null;

   function Iter_Next
     (Container : Multiset; Key : Private_Key) return Private_Key
   with
     Ghost  => SPARKlib_Full,
     Global => null,
     Pre    => Iter_Has_Element (Container, Key);

   function Iter_Element
     (Container : Multiset; Key : Private_Key) return Element_Type
   with
     Ghost  => SPARKlib_Full,
     Global => null,
     Pre    => Iter_Has_Element (Container, Key);

   ------------------------------------------
   -- Additional Primitives For Aggregates --
   ------------------------------------------

   function Aggr_Eq_Elements (Left, Right : Element_Type) return Boolean
   is (Equivalent_Elements (Left, Right))
   with
     Global   => null,
     Ghost    => Static,
     Annotate => (GNATprove, Inline_For_Proof),
     Annotate => (GNATprove, Container_Aggregates, "Equivalent_Keys");

   function Default_Count return Big_Natural
   is (0)
   with
     Ghost    => Static,
     Annotate => (GNATprove, Inline_For_Proof),
     Annotate => (GNATprove, Container_Aggregates, "Default_Item");

   procedure Aggr_Include
     (Container : in out Multiset; Element : Element_Type; Count : Big_Natural)
   with
     Global => null,
     Pre    => (SPARKlib_Defensive => Nb_Occurence (Container, Element) = 0),
     Post   =>
       (SPARKlib_Full =>
          Nb_Occurence (Container, Element) = Count
          and then
            (for all E of Container =>
               (if not Equivalent_Elements (E, Element)
                then
                  Nb_Occurence (Container, E)
                  = Nb_Occurence (Container'Old, E)))
          and then
            (for all E of Container'Old =>
               (if not Equivalent_Elements (E, Element)
                then
                  Nb_Occurence (Container, E)
                  = Nb_Occurence (Container'Old, E))));

private
   pragma SPARK_Mode (Off); --  #BODYMODE
   package Maps is new
     SPARK.Containers.Functional.Maps
       (Key_Type                   => Element_Type,
        Element_Type               => Big_Positive,
        Equivalent_Keys            => Equivalent_Elements,
        Equivalent_Keys_Reflexive  => Eq_Elements_Checks.Eq_Reflexive,
        Equivalent_Keys_Symmetric  => Eq_Elements_Checks.Eq_Symmetric,
        Equivalent_Keys_Transitive => Eq_Elements_Checks.Eq_Transitive);
   use Maps;

   type Multiset is record
      Map  : Maps.Map;
      Card : Big_Natural := 0;
   end record
   with Type_Invariant => (SPARKlib_Full => Invariant (Map, Card));

   function Invariant (Container : Map; Card : Big_Natural) return Boolean
   with Ghost => SPARKlib_Full;

   ---------------------------------------
   -- Iteration on Functional Multisets --
   ---------------------------------------

   type Iterable_Multiset is record
      Map  : Maps.Iterable_Map := Iterate (Empty_Map);
      Card : Big_Natural := 0;
   end record
   with Type_Invariant => (SPARKlib_Full => Invariant (Get_Map (Map), Card));

   function Element
     (Iterator : Iterable_Multiset; Cursor : Multiset) return Element_Type
   is (Choose (Cursor));

   function First (Iterator : Iterable_Multiset) return Multiset
   is (Map => First (Iterator.Map), Card => Iterator.Card);

   function Get_Multiset (Iterator : Iterable_Multiset) return Multiset
   is (Map => First (Iterator.Map), Card => Iterator.Card);

   function Has_Element
     (Iterator : Iterable_Multiset; Cursor : Multiset) return Boolean
   is (Valid_Subset (Iterator, Cursor) and then Length (Cursor.Map) > 0);

   function Iterate (Container : Multiset) return Iterable_Multiset
   is (Map => Iterate (Container.Map), Card => Container.Card);

   function Multiset_Logic_Equal (Left, Right : Multiset) return Boolean
   is (Map_Logic_Equal (Left.Map, Right.Map));

   function Valid_Subset
     (Iterator : Iterable_Multiset; Cursor : Multiset) return Boolean
   is (Valid_Submap (Iterator.Map, Cursor.Map));

   --------------------------------------------------
   -- Iteration Primitives Used For Quantification --
   --------------------------------------------------

   type Private_Key is new Maps.Private_Key;

   function Iter_First (Container : Multiset) return Private_Key
   is (Iter_First (Container.Map));

   function Iter_Has_Element
     (Container : Multiset; Key : Private_Key) return Boolean
   is (Iter_Has_Element (Container.Map, Key));

   function Iter_Next
     (Container : Multiset; Key : Private_Key) return Private_Key
   is (Iter_Next (Container.Map, Key));

   function Iter_Element
     (Container : Multiset; Key : Private_Key) return Element_Type
   is (Iter_Element (Container.Map, Key));

end SPARK.Containers.Functional.Multisets;
