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

import Vector
import Util
import FIFO
import FIFOF
import SpecialFIFOs

#define UIntLog2N(n) (UInt (TLog n))

interface (TagEngine :: # -> *) numTags =
    requestTag :: ActionValue UIntLog2N(numTags)
    retireTag  :: UIntLog2N(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)

    -- State
    tagUsage :: Vector numTags (Reg Bool) <- replicateM (mkReg False)  -- Tracks which tags are in use
    initialTagCounter <- mkReg (Just (maxTagCount - 1))  -- Distributes initial tags
    retireQueue <- mkBypassFIFO  -- Queue for tags being retired
    freeTagQueue <- mkSizedFIFOF 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 UIntLog2N(numTags)
            requestTag = do
                case initialTagCounter of
                    Just 0 -> do
                        initialTagCounter := Nothing
                        requestSignal.wset 0
                        return 0
                    Just tag -> do
                        initialTagCounter := Just (tag - 1)
                        requestSignal.wset tag
                        return tag
                    Nothing -> do
                        let tag = freeTagQueue.first
                        freeTagQueue.deq
                        requestSignal.wset tag
                        return tag

            retireTag :: UIntLog2N(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 {}