package TagEngine(
    MkTagType,
    TagEngine(..),
    Util.BasicResult(..),
    mkTagEngine) where

import Vector
import Util
import FIFO
import SpecialFIFOs

type MkTagType numTags = (UInt (TLog (TAdd 1 numTags)))

interface (TagEngine :: # -> *) numTags =
    requestTag :: ActionValue (MkTagType numTags)
    retireTag  :: (MkTagType numTags) -> Action

-- The tag engine returns a tag that is unique for the duration of
-- the lifetime of the tag. Useful when you need to tag transactions
-- on a bus for example.
-- This implementation is FIFO based.
mkTagEngine :: Module (TagEngine numTags)
mkTagEngine = do
    -- Constants
    let maxTagCount = fromInteger (valueOf numTags)

    tagUsage :: Vector numTags (Reg Bool) <- replicateM (mkReg False)  -- Tracks which tags are in use

    -- Since Bluespec doesn't allow us to initialize FIFOs with values at
    -- reset, we can pretend as if the buffer within our freeTagQueue is populated
    -- with sequentially incrementing values(starting from 0) on reset
    -- by having our tag engine effectively return the value of a decrementing
    -- counter initialized to (maxTagCount - 1) for the first n tag requests made
    -- to TagEngine where `n := maxTagCount`.
    initialTagDistributor <- mkReg (Just (maxTagCount - 1))  -- Distributes initial tags
    retireQueue <- mkBypassFIFO  -- Queue for tags being retired
    freeTagQueue <- mkSizedFIFO maxTagCount  -- Queue of available tags

    -- Signals
    retireSignal <- mkRWire  -- Signals a tag retirement
    requestSignal <- mkRWire  -- Signals a tag request

    -- Debug
    debugOnce <- mkReg True

    -- Rules
    addRules |>
        rules
            "debug_initial_state": when debugOnce ==> do
                $display "tagUsage: " (fshow (readVReg tagUsage))
                debugOnce := False

            "retire_tag": when True ==> do
                let tag = retireQueue.first
                $display "Retiring tag: " (fshow tag)
                retireQueue.deq
                freeTagQueue.enq tag
                retireSignal.wset tag

            -- Combined update rules (simplified below)
            "update_usage": when True ==> do
                let mRetireTag = retireSignal.wget
                    mRequestTag = requestSignal.wget
                case (mRetireTag, mRequestTag) of
                    (Just retireTag, Just requestTag) -> do
                        let usage = readVReg tagUsage
                            usage' = update usage requestTag True
                            usage'' = update usage' retireTag False
                        writeVReg tagUsage usage''
                        $display $time " Updated usage (request + retire): " (fshow |> readVReg tagUsage)
                    (Just retireTag, Nothing) -> do
                        (select tagUsage retireTag) := False
                        $display $time " Updated usage (retire): " (fshow (readVReg tagUsage))
                    (Nothing, Just requestTag) -> do
                        (select tagUsage requestTag) := True
                        $display $time " Updated usage (request): " (fshow (readVReg tagUsage))
                    (Nothing, Nothing) -> action {}

    -- Interface
    return |>
        interface TagEngine
            requestTag :: ActionValue (MkTagType numTags)
            requestTag = do
                case initialTagDistributor of
                    Just 0 -> do
                        initialTagDistributor := Nothing
                        requestSignal.wset 0
                        return 0
                    Just tag -> do
                        initialTagDistributor := Just (tag - 1)
                        requestSignal.wset tag
                        return tag
                    Nothing -> do
                        let tag = freeTagQueue.first
                        freeTagQueue.deq
                        requestSignal.wset tag
                        return tag

            -- `retireTag` isn't guarded on tag validity(this would break Bluespec's safety model)
            -- so it is advisable that the caller of `retireTag` only attempt to retire valid tags.
            -- Internally, the tagEngine will keep a correct and consistent state since TagEngine
            -- validates tags before attempting to retire them.
            retireTag :: (MkTagType numTags) -> Action
            retireTag tag =
                do
                    let
                        tagValid = tag < maxTagCount
                        tagInUse = readReg (select tagUsage tag)
                    if (tagValid && tagInUse)
                        then do
                            retireQueue.enq tag
                        else do
                            action {}