    module StackPlacements
      ( SlotSet, allStackSlots  -- the infinite set of stack slots
      , SlotClass(..), slotClassBits, stackSlot32, stackSlot64, stackSlot128
      , allSlotClasses
      , getStackSlot, extendSlotSet, deleteFromSlotSet, elemSlotSet, chooseSlot
      , StackPlacement(..)
      )
    where

    import Maybes
    import Outputable
    import Unique

    import Prelude hiding (pi)
    import Data.List

    {-

    The goal here is to provide placements on the stack that will allow,
    for example, two 32-bit words to spill to a slot previously used by a
    64-bit floating-point value.  I use a simple buddy-system allocator
    that splits large slots in half as needed; this will work fine until
    the day when somebody wants to spill an 80-bit Intel floating-point
    register into the Intel standard 96-bit stack slot.

    -}

    data SlotClass = SlotClass32 | SlotClass64 | SlotClass128
      deriving (Eq)

    instance Uniquable SlotClass where
        getUnique = getUnique . slotClassBits

    instance Outputable SlotClass where
        ppr cls = text "class of" <+> int (slotClassBits cls) <> text "-bit stack slots"

    slotClassBits :: SlotClass -> Int
    slotClassBits SlotClass32 = 32
    slotClassBits SlotClass64 = 64
    slotClassBits SlotClass128 = 128

    data StackPlacement = FullSlot SlotClass Int
                   | YoungHalf StackPlacement
                   | OldHalf StackPlacement
      deriving (Eq)

    data OneSize = OneSize { full_slots :: [StackPlacement], fragments :: [StackPlacement] }
      -- ^ Always used for slots that have been previously used

    data SlotSet = SlotSet { s32, s64, s128 :: OneSize, next_unused :: Int }

    allStackSlots :: SlotSet
    allStackSlots = SlotSet empty empty empty 0
        where empty = OneSize [] []

    psize :: StackPlacement -> Int
    psize (FullSlot cls _) = slotClassBits cls
    psize (YoungHalf p) = psize p `div` 2
    psize (OldHalf   p) = psize p `div` 2

    -- | Get a slot no matter what
    get32, get64, get128 :: SlotSet -> (StackPlacement, SlotSet)

    -- | Get a previously used slot if one exists
    getu32, getu64, getu128 :: SlotSet -> Maybe (StackPlacement, SlotSet)

    -- | Only supported slot classes

    stackSlot32, stackSlot64, stackSlot128 :: SlotClass
    stackSlot32  = SlotClass32
    stackSlot64  = SlotClass64
    stackSlot128 = SlotClass128

    allSlotClasses :: [SlotClass]
    allSlotClasses = [stackSlot32, stackSlot64, stackSlot128]

    -- | Get a fresh slot, never before used
    getFull :: SlotClass -> SlotSet -> (StackPlacement, SlotSet)

    infixr 4 |||

    (|||) :: (SlotSet -> Maybe (StackPlacement, SlotSet)) ->
             (SlotSet ->       (StackPlacement, SlotSet)) ->
             (SlotSet ->       (StackPlacement, SlotSet))

    f1 ||| f2 = \slots -> f1 slots `orElse`   f2 slots

    getFull cls slots = (FullSlot cls n, slots { next_unused = n + 1 })
        where n = next_unused slots

    get32  = getu32  ||| (fmap split64  . getu64)  ||| getFull stackSlot32
    get64  = getu64  ||| (fmap split128 . getu128) ||| getFull stackSlot64
    get128 = getu128 ||| getFull stackSlot128

    type SizeGetter = SlotSet -> OneSize
    type SizeSetter = OneSize -> SlotSet -> SlotSet

    upd32, upd64, upd128 :: SizeSetter
    upd32  this_size slots = slots { s32  = this_size }
    upd64  this_size slots = slots { s64  = this_size }
    upd128 this_size slots = slots { s128 = this_size }

    with_size :: Int -> (SizeGetter -> SizeSetter -> a) -> a
    with_size  32 = with_32
    with_size  64 = with_64
    with_size 128 = with_128
    with_size _   = panic "non-standard slot size -- error in size computation?"

    with_32, with_64, with_128 :: (SizeGetter -> SizeSetter -> a) -> a
    with_32  f = f s32  upd32
    with_64  f = f s64  upd64
    with_128 f = f s128 upd128

    getu32  = with_32  getUsed
    getu64  = with_64  getUsed
    getu128 = with_128 getUsed

    getUsed :: SizeGetter -> SizeSetter -> SlotSet -> Maybe (StackPlacement, SlotSet)
    getUsed get set slots =
        let this_size = get slots in
        case full_slots this_size of
          p : ps -> Just (p, set (this_size { full_slots = ps }) slots)
          [] -> case fragments this_size of
                  p : ps -> Just (p, set (this_size { fragments = ps }) slots)
                  [] -> Nothing

    -- | When splitting, allocate the old half first in case it makes the
    -- stack smaller at a call site.
    split64, split128 :: (StackPlacement, SlotSet) -> (StackPlacement, SlotSet)
    split64  (p, slots) = (OldHalf p, slots { s32 = cons_frag (YoungHalf p) (s32 slots) })
    split128 (p, slots) = (OldHalf p, slots { s64 = cons_frag (YoungHalf p) (s64 slots) })

    cons_frag :: StackPlacement -> OneSize -> OneSize
    cons_frag p this_size = this_size { fragments = p : fragments this_size }

    ----------------------------
    instance Outputable StackPlacement where
      ppr (FullSlot cls n) = int (slotClassBits cls) <> text "-bit slot " <> int n
      ppr (YoungHalf p) = text "young half of" <+> ppr p
      ppr (OldHalf   p) = text "old half of"   <+> ppr p

    instance Outputable SlotSet where
      ppr slots = fsep $ punctuate comma
                  (pprSlots (s32 slots) ++ pprSlots (s64 slots) ++ pprSlots (s128 slots) ++
                   [text "and slots numbered" <+> int (next_unused slots)
                             <+> text "and up"])
       where pprSlots (OneSize w fs) = map ppr w ++ map ppr fs

    {-
    instance ColorSet SlotSet SlotClass StackPlacement where
      emptyColorSet = panic "The set of stack slots is never empty"
      deleteFromColorSet = deleteFromSlotSet
      extendColorSet slots (cls, p@(FullSlot {})) =
          with_size (slotClassBits cls) add_full p (pi slots)
      extendColorSet slots (cls, p) = with_size (slotClassBits cls) add_frag p (pi slots)
      chooseColor        = chooseSlot
    -}

    deleteFromSlotSet :: StackPlacement -> SlotSet -> SlotSet
    deleteFromSlotSet p@(FullSlot {}) slots = with_size (psize p) remove_full p (pi slots)
    deleteFromSlotSet p               slots = with_size (psize p) remove_frag p (pi slots)

    extendSlotSet :: SlotSet -> StackPlacement -> SlotSet
    extendSlotSet slots p@(FullSlot {}) = with_size (psize p) add_full p (pi slots)
    extendSlotSet slots p               = with_size (psize p) add_frag p (pi slots)

    elemSlotSet :: StackPlacement -> SlotSet -> Bool
    elemSlotSet p@(FullSlot {}) slots = with_size (psize p) elem_full p slots
    elemSlotSet p               slots = with_size (psize p) elem_frag p slots

    remove_full, remove_frag, add_full, add_frag
        :: SizeGetter -> SizeSetter -> StackPlacement -> SlotSet -> SlotSet

    remove_full get set p slots = set p' slots
        where this_size = get slots
              p' = this_size { full_slots = delete p $ full_slots this_size }

    remove_frag get set p slots = set p' slots
        where this_size = get slots
              p' = this_size { full_slots = delete p $ full_slots this_size }

    add_full get set p slots = set p' slots
        where this_size = get slots
              p' = this_size { full_slots = add p $ full_slots this_size }

    add_frag get set p slots = set p' slots
        where this_size = get slots
              p' = this_size { full_slots = add p $ full_slots this_size }

    add :: Eq a => a -> [a] -> [a]
    add x xs = if notElem x xs then x : xs else xs

    elem_full, elem_frag :: SizeGetter -> SizeSetter -> StackPlacement -> SlotSet -> Bool
    elem_full get _set p slots = elem p (full_slots $ get slots)
    elem_frag get _set p slots = elem p (fragments  $ get slots)

    getStackSlot :: SlotClass -> SlotSet -> (StackPlacement, SlotSet)
    getStackSlot cls slots =
      case cls of
        SlotClass32  -> get32  (pi slots)
        SlotClass64  -> get64  (pi slots)
        SlotClass128 -> get128 (pi slots)

    chooseSlot :: SlotClass -> [StackPlacement] -> SlotSet -> Maybe (StackPlacement, SlotSet)
    chooseSlot cls prefs slots =
      case filter (flip elemSlotSet slots) prefs of
        placement : _ -> Just (placement, deleteFromSlotSet placement (pi slots))
        [] -> Just (getStackSlot cls slots)

    check_invariant :: Bool
    check_invariant = True

    pi :: SlotSet -> SlotSet
    pi = if check_invariant then panic_on_invariant_violation else id

    panic_on_invariant_violation :: SlotSet -> SlotSet
    panic_on_invariant_violation slots =
        check 32 (s32 slots) $ check 64 (s64 slots) $ check 128 (s128 slots) $ slots
      where n = next_unused slots
            check bits this_size = (check_full bits $ full_slots this_size) .
                                   (check_frag bits $ fragments  this_size)
            check_full _ [] = id
            check_full bits (FullSlot cls k : ps) =
                if slotClassBits cls /= bits then panic "slot in bin of wrong size"
                else if k >= n then panic "slot number is unreasonably fresh"
                     else check_full bits ps
            check_full _ _ = panic "a fragment is in a bin reserved for full slots"
            check_frag _ [] = id
            check_frag _ (FullSlot {} : _) =
                panic "a full slot is in a bin reserved for fragments"
            check_frag bits (p : ps) =
                if bits /= psize p then panic "slot in bin of wrong size"
                else if pnumber p >= n then panic "slot number is unreasonably fresh"
                     else check_frag bits ps
            pnumber (FullSlot _ k) = k
            pnumber (YoungHalf p) = pnumber p
            pnumber (OldHalf p)   = pnumber p