/*
 * Copyright 2010, Ingo Weinhold, ingo_weinhold@gmx.de.
 * Distributed under the terms of the MIT License.
 */


#include "BlockAllocator.h"

#include <string.h>

#include <algorithm>

#include <util/AutoLock.h>

#include "Block.h"
#include "checksumfs.h"
#include "DebugSupport.h"
#include "SuperBlock.h"
#include "Volume.h"


// This is a simple block allocator implementation. It manages the on-disk
// block bitmap structures. Those consist of:
// * The block bitmap itself: A contiguous run of blocks with one bit for each
//   of the blocks accessible by the file system, indicating whether the block
//   is used or free. The bit is set for a used, clear for a free block. For
//   faster access the bits are grouped in 32 bit integers (host endian ATM).
// * The allocation groups blocks: An allocation group consists of the blocks
//   that 2048 block bitmap blocks refer to. Each allocation group is
//   represented by a single block which contains 2048 16 bit numbers (host
//   endian ATM), each of which referring to a block bitmap block and encoding
//   how many free blocks the block bitmap block refers to.
// The blocks for the allocation groups are directly followed by the block
// bitmap blocks. The blocks beyond the blocks accessible by the file system
// that the last block bitmap block may refer to are marked used. Non-existing
// block bitmap blocks referred to by the last allocation group block are set
// to have 0 free blocks.


static const uint32 kBlocksPerBitmapBlock       = B_PAGE_SIZE * 8;
static const uint32 kBitmapBlocksPerGroup       = B_PAGE_SIZE / 2;
static const uint32 kBlocksPerGroup
        = kBitmapBlocksPerGroup * kBlocksPerBitmapBlock;


BlockAllocator::BlockAllocator(Volume* volume)
        :
        fVolume(volume),
        fTotalBlocks(volume->TotalBlocks()),
        fBitmapBlockCount((fTotalBlocks + B_PAGE_SIZE - 1) / B_PAGE_SIZE)
{
        fAllocationGroupCount = (fBitmapBlockCount + kBitmapBlocksPerGroup - 1)
                / kBitmapBlocksPerGroup;
        mutex_init(&fLock, "checksumfs block allocator");
}


BlockAllocator::~BlockAllocator()
{
        mutex_destroy(&fLock);
}


status_t
BlockAllocator::Init(uint64 blockBitmap, uint64 freeBlocks)
{
        if (blockBitmap <= kCheckSumFSSuperBlockOffset / B_PAGE_SIZE
                || blockBitmap >= fTotalBlocks
                || fTotalBlocks - blockBitmap < fAllocationGroupCount) {
                return B_BAD_DATA;
        }

        fFreeBlocks = freeBlocks;
        fAllocationGroupBlock = blockBitmap;
        fBitmapBlock = blockBitmap + fAllocationGroupCount;

        return B_OK;
}


status_t
BlockAllocator::Initialize(Transaction& transaction)
{
        status_t error = Init(kCheckSumFSSuperBlockOffset / B_PAGE_SIZE + 1,
                fTotalBlocks);
        if (error != B_OK)
                return error;

        PRINT("BlockAllocator::Initialize():\n");
        PRINT("  fTotalBlocks:          %" B_PRIu64 "\n", fTotalBlocks);
        PRINT("  fFreeBlocks:           %" B_PRIu64 "\n", fFreeBlocks);
        PRINT("  fAllocationGroupBlock: %" B_PRIu64 "\n", fAllocationGroupBlock);
        PRINT("  fAllocationGroupCount: %" B_PRIu64 "\n", fAllocationGroupCount);
        PRINT("  fBitmapBlock:          %" B_PRIu64 "\n", fBitmapBlock);
        PRINT("  fBitmapBlockCount:     %" B_PRIu64 "\n", fBitmapBlockCount);

        // clear the block bitmap
        for (uint64 i = 0; i < fBitmapBlockCount; i++) {
                Block block;
                if (!block.GetZero(fVolume, fBitmapBlock + i, transaction))
                        return B_ERROR;
        }

        // the last block of the block bitmap may be partial -- mark the blocks
        // beyond the end used
        uint32 partialBitmapBlock = fTotalBlocks % kBlocksPerBitmapBlock;
        if (partialBitmapBlock != 0) {
                Block block;
                if (!block.GetZero(fVolume, fBitmapBlock + fBitmapBlockCount - 1,
                                transaction)) {
                        return B_ERROR;
                }

                // set full uint32s
                uint32* bits = (uint32*)block.Data();
                uint32 offset = (partialBitmapBlock + 31) / 32;
                if (partialBitmapBlock <= B_PAGE_SIZE - 4)
                        memset(bits + offset, 0xff, B_PAGE_SIZE - offset * 4);

                // set the partial uint32
                uint32 bitOffset = partialBitmapBlock % 32;
                if (bitOffset != 0)
                        bits[offset] = ~(((uint32)1 << bitOffset) - 1);
        }

        // init the allocation groups
        uint32 partialGroup = fTotalBlocks % kBlocksPerGroup;
        for (uint64 i = 0; i < fAllocationGroupCount; i++) {
                Block block;
                if (!block.GetZero(fVolume, fAllocationGroupBlock + i, transaction))
                        return B_ERROR;

                uint16* counts = (uint16*)block.Data();

                if (i < fAllocationGroupCount - 1 || partialGroup == 0) {
                        // not the last group or the last group is not partial -- all
                        // blocks are free
                        for (uint32 i = 0; i < kBitmapBlocksPerGroup; i++)
                                counts[i] = kBlocksPerBitmapBlock;
                } else {
                        // the last, partial group
                        if (partialGroup != 0) {
                                uint32 offset = partialGroup / kBlocksPerBitmapBlock;
                                for (uint32 i = 0; i < offset; i++)
                                        counts[i] = kBlocksPerBitmapBlock;
                                counts[offset] = partialBitmapBlock;
                        }
                }
        }

        // mark all blocks we already use used
        error = AllocateExactly(0, fBitmapBlock + fBitmapBlockCount, transaction);
        if (error != B_OK)
                return error;

        PRINT("BlockAllocator::Initialize() done:\n");
        PRINT("  fFreeBlocks:           %" B_PRIu64 "\n", fFreeBlocks);

        return B_OK;
}


status_t
BlockAllocator::Allocate(uint64 baseHint, uint64 count,
        Transaction& transaction, uint64& _allocatedBase, uint64& _allocatedCount)
{
        MutexLocker locker(fLock);

        PRINT("BlockAllocator::Allocate(%" B_PRIu64 ", %" B_PRIu64 ")\n", baseHint,
                count);

        if (fFreeBlocks == 0)
                return B_DEVICE_FULL;

        if (baseHint >= fTotalBlocks)
                baseHint = 0;

        // search from base hint to end
        status_t error = _Allocate(baseHint, fTotalBlocks, count, transaction,
                &_allocatedBase, _allocatedCount);
        if (error == B_OK)
                return _UpdateSuperBlock(transaction);
        if (baseHint == 0)
                return error;

        // search from 0 to hint
        error = _Allocate(0, baseHint, count, transaction, &_allocatedBase,
                _allocatedCount);
        if (error != B_OK)
                return error;

        return _UpdateSuperBlock(transaction);
}


status_t
BlockAllocator::AllocateExactly(uint64 base, uint64 count,
        Transaction& transaction)
{
        MutexLocker locker(fLock);

        PRINT("BlockAllocator::AllocateExactly(%" B_PRIu64 ", %" B_PRIu64 ")\n",
                base, count);

        uint64 allocated;
        status_t error = _Allocate(base, fTotalBlocks, count, transaction, NULL,
                allocated);
        if (error != B_OK)
                return error;

        if (allocated < count)
                return B_BUSY;

        return _UpdateSuperBlock(transaction);
}


status_t
BlockAllocator::Free(uint64 base, uint64 count, Transaction& transaction)
{
        MutexLocker locker(fLock);

        status_t error = _Free(base, count, transaction);
        if (error != B_OK)
                return error;

        return _UpdateSuperBlock(transaction);
}


void
BlockAllocator::ResetFreeBlocks(uint64 count)
{
        MutexLocker locker(fLock);

        fFreeBlocks = count;
}


/*!     Allocates contiguous blocks.

        Might allocate fewer block than requested. If no block could be allocated
        at all, an error is returned.
        If \a _allocatedBase is not \c NULL, the method may move the base up, if it
        isn't able to allocate anything at the given base.

        \param base The first potential block to allocate.
        \param searchEnd A hint for the method where to stop searching for free
                blocks.
        \param count The maximum number of blocks to allocate.
        \param _allocatedBase If not \c NULL, the may allocate at a greater base.
                The base of the actual allocation is returned via this variable.
        \param _allocatedCount On success the variable will be set to the number of
                blocks actually allocated.
        \return \c B_OK, if one or more blocks could be allocated, another error
                code otherwise.
*/
status_t
BlockAllocator::_Allocate(uint64 base, uint64 searchEnd, uint64 count,
        Transaction& transaction, uint64* _allocatedBase, uint64& _allocatedCount)
{
        ASSERT(base <= fTotalBlocks);
        ASSERT(searchEnd <= fTotalBlocks);

        if (base >= searchEnd || fFreeBlocks == 0)
                RETURN_ERROR(B_BUSY);

        uint64 groupOffset = base % kBlocksPerGroup;
        uint64 remaining = count;

        // If we're allowed to move the base, loop until we allocate something.
        if (_allocatedBase != NULL) {
                while (count > 0 && base < searchEnd) {
                        uint64 toAllocate = std::min(count, kBlocksPerGroup - groupOffset);

                        uint32 allocated;
                        status_t error = _AllocateInGroup(base, searchEnd, toAllocate,
                                transaction, _allocatedBase, allocated);
                        if (error == B_OK) {
                                fFreeBlocks -= toAllocate;

                                if (allocated == toAllocate) {
                                        _allocatedCount = allocated;
                                        return B_OK;
                                }

                                count = std::min(count,
                                        searchEnd - (uint32)*_allocatedBase % kBlocksPerGroup);
                                _allocatedBase = NULL;
                                remaining = count - allocated;

                                break;
                        }

                        // nothing yet -- continue with the next group
                        count -= toAllocate;
                        base += toAllocate;
                        groupOffset = 0;
                }

                if (_allocatedBase != NULL)
                        return B_BUSY;
        }

        // We're not/no longer allowed to move the base. Loop as long as we can
        // allocate what ask for.
        while (remaining > 0 && base < searchEnd) {
                uint64 toAllocate = std::min(remaining, kBlocksPerGroup - groupOffset);
                uint32 allocated;
                status_t error = _AllocateInGroup(base, searchEnd, toAllocate,
                        transaction, NULL, allocated);
                if (error != B_OK)
                        break;

                fFreeBlocks -= toAllocate;
                remaining -= allocated;

                if (allocated < toAllocate)
                        break;

                base += toAllocate;
                groupOffset = 0;
        }

        if (remaining == count)
                RETURN_ERROR(B_BUSY);

        _allocatedCount = count - remaining;
        return B_OK;
}


/*!     Allocates contiguous blocks in an allocation group.

        The range specified by \a base and \a count must lie fully within a single
        allocation group.

        Might allocate fewer block than requested. If no block could be allocated
        at all, an error is returned.
        If \a _allocatedBase is not \c NULL, the method may move the base up, if it
        isn't able to allocate anything at the given base.

        \param base The first potential block to allocate.
        \param searchEnd A hint for the method where to stop searching for free
                blocks.
        \param count The maximum number of blocks to allocate.
        \param _allocatedBase If not \c NULL, the may allocate at a greater base.
                The base of the actual allocation is returned via this variable.
        \param _allocatedCount On success the variable will be set to the number of
                blocks actually allocated.
        \return \c B_OK, if one or more blocks could be allocated, another error
                code otherwise.
*/
status_t
BlockAllocator::_AllocateInGroup(uint64 base, uint64 searchEnd, uint32 count,
        Transaction& transaction, uint64* _allocatedBase, uint32& _allocatedCount)
{
        PRINT("BlockAllocator::_AllocateInGroup(%" B_PRIu64 ", %" B_PRIu32 ")\n",
                base, count);

        ASSERT(count <= kBlocksPerGroup);
        ASSERT(base % kBlocksPerGroup + count <= kBlocksPerGroup);

        if (base >= searchEnd)
                RETURN_ERROR(B_BUSY);

        Block block;
        if (!block.GetWritable(fVolume,
                        fAllocationGroupBlock + base / kBlocksPerGroup, transaction)) {
                RETURN_ERROR(B_ERROR);
        }

        uint16* counts = (uint16*)block.Data();

        uint32 blockIndex = base / kBlocksPerBitmapBlock % kBitmapBlocksPerGroup;
        uint64 inBlockOffset = base % kBlocksPerBitmapBlock;
        uint64 remaining = count;

        // If we're allowed to move the base, skip used blocks.
        if (_allocatedBase != NULL) {
                // check partial block
                if (inBlockOffset != 0) {
                        if (counts[blockIndex] > 0) {
                                uint32 allocated;
                                if (_AllocateInBitmapBlock(base, count, transaction,
                                                _allocatedBase, allocated) == B_OK) {
                                        counts[blockIndex] -= allocated;

                                        if (inBlockOffset + allocated < kBlocksPerBitmapBlock
                                                || allocated == remaining) {
                                                _allocatedCount = allocated;
                                                return B_OK;
                                        }

                                        count = std::min(count,
                                                kBlocksPerGroup
                                                        - (uint32)*_allocatedBase % kBlocksPerGroup);
                                        _allocatedBase = NULL;
                                        remaining = count - allocated;
                                }
                        }

                        base += kBlocksPerBitmapBlock - inBlockOffset;
                        inBlockOffset = 0;
                        blockIndex++;
                }

                // skip completely used blocks
                if (_allocatedBase != NULL) {
                        while (blockIndex < kBitmapBlocksPerGroup && base < searchEnd) {
                                if (counts[blockIndex] > 0)
                                        break;

                                base += kBlocksPerBitmapBlock;
                                blockIndex++;
                        }
                }

                if (blockIndex == kBitmapBlocksPerGroup)
                        return B_BUSY;

                // Clamp the count to allocate, if we have moved the base too far. Do
                // this only, if we haven't allocated anything so far.
                if (_allocatedBase != NULL) {
                        count = std::min(count,
                                kBlocksPerGroup - (uint32)base % kBlocksPerGroup);
                        remaining = count;
                }
        }

        // Allocate as many of the requested blocks as we can.
        while (remaining > 0 && base < searchEnd) {
                if (counts[blockIndex] == 0)
                        break;

                uint32 toAllocate = std::min(remaining,
                        kBlocksPerBitmapBlock - inBlockOffset);

                uint32 allocated;
                status_t error = _AllocateInBitmapBlock(base, toAllocate, transaction,
                        _allocatedBase, allocated);
                if (error != B_OK)
                        break;

                counts[blockIndex] -= allocated;
                remaining -= allocated;

                if (allocated < toAllocate)
                        break;

                base += allocated;
                blockIndex++;
                inBlockOffset = 0;
                _allocatedBase = NULL;
        }

        if (remaining == count)
                return B_BUSY;

        _allocatedCount = count - remaining;
        return B_OK;
}


/*!     Allocates contiguous blocks in a bitmap block.

        The range specified by \a base and \a count must lie fully within a single
        bitmap block.

        Might allocate fewer block than requested. If no block could be allocated
        at all, an error is returned.
        If \a _allocatedBase is not \c NULL, the method may move the base up, if it
        isn't able to allocate anything at the given base.

        \param base The first potential block to allocate.
        \param count The maximum number of blocks to allocate.
        \param _allocatedBase If not \c NULL, the may allocate at a greater base.
                The base of the actual allocation is returned via this variable.
        \param _allocatedCount On success the variable will be set to the number of
                blocks actually allocated.
        \return \c B_OK, if one or more blocks could be allocated, another error
                code otherwise.
*/
status_t
BlockAllocator::_AllocateInBitmapBlock(uint64 base, uint32 count,
        Transaction& transaction, uint64* _allocatedBase, uint32& _allocatedCount)
{
        PRINT("BlockAllocator::_AllocateInBitmapBlock(%" B_PRIu64 ", %" B_PRIu32
                ")\n", base, count);

        ASSERT(count <= kBlocksPerBitmapBlock);
        ASSERT(base % kBlocksPerBitmapBlock + count <= kBlocksPerBitmapBlock);

        Block block;
        if (!block.GetWritable(fVolume,
                        fBitmapBlock + base / kBlocksPerBitmapBlock, transaction)) {
                RETURN_ERROR(B_ERROR);
        }

        uint32* bits = (uint32*)block.Data()
                + base % kBlocksPerBitmapBlock / 32;
        uint32* const bitsEnd = (uint32*)block.Data() + kBlocksPerBitmapBlock / 32;
        uint32 bitOffset = base % 32;

        // If we're allowed to move the base, skip used blocks.
        if (_allocatedBase != NULL) {
                // check partial uint32 at the beginning
                bool foundBase = false;
                if (bitOffset > 0) {
                        uint32 mask = ~(((uint32)1 << bitOffset) - 1);
                        if ((*bits & mask) != mask) {
                                while ((*bits & ((uint32)1 << bitOffset)) != 0) {
                                        bitOffset++;
                                        base++;
                                }
                                foundBase = true;
                        } else {
                                // all used -- skip
                                bits++;
                                base += 32 - bitOffset;
                                bitOffset = 0;
                        }
                }

                // check complete uint32s
                if (!foundBase) {
                        while (bits < bitsEnd) {
                                if (*bits != 0xffffffff) {
                                        bitOffset = 0;
                                        while ((*bits & ((uint32)1 << bitOffset)) != 0) {
                                                bitOffset++;
                                                base++;
                                        }
                                        foundBase = true;
                                        break;
                                }

                                bits++;
                                base += 32;
                        }
                }

                if (!foundBase)
                        return B_BUSY;

                // Clamp the count to allocate, if we have moved the base too far.
                if (base % kBlocksPerBitmapBlock + count > kBlocksPerBitmapBlock)
                        count = kBlocksPerBitmapBlock - base % kBlocksPerBitmapBlock;
        }

        // Allocate as many of the requested blocks as we can, starting at the base
        // we have.
        uint32 remaining = count;

        while (remaining > 0 && bits < bitsEnd) {
                PRINT("  remaining: %" B_PRIu32 ", index: %" B_PRIu32 ".%" B_PRIu32
                        ", bits: %#" B_PRIx32 "\n", remaining,
                        kBlocksPerBitmapBlock - (bitsEnd - bits), bitOffset, *bits);

                // TODO: Not particularly efficient for large allocations.
                uint32 endOffset = std::min(bitOffset + remaining, (uint32)32);
                for (; bitOffset < endOffset; bitOffset++) {
                        if ((*bits & ((uint32)1 << bitOffset)) != 0) {
                                bits = bitsEnd;
                                break;
                        }

                        *bits |= (uint32)1 << bitOffset;
                        remaining--;
                }

                bits++;
                bitOffset = 0;
        }

        if (remaining == count)
                RETURN_ERROR(B_BUSY);

        _allocatedCount = count - remaining;
        if (_allocatedBase != NULL)
                *_allocatedBase = base;

        return B_OK;
}


status_t
BlockAllocator::_Free(uint64 base, uint64 count, Transaction& transaction)
{
        if (count == 0)
                return B_OK;

        PRINT("BlockAllocator::_Free(%" B_PRIu64 ", %" B_PRIu64 ")\n", base, count);

        ASSERT(count <= fTotalBlocks - fFreeBlocks);
        ASSERT(base < fTotalBlocks && fTotalBlocks - base >= count);

        uint64 groupOffset = base % kBlocksPerGroup;
        uint64 remaining = count;

        while (remaining > 0) {
                uint64 toFree = std::min(remaining, kBlocksPerGroup - groupOffset);
                status_t error = _FreeInGroup(base, toFree, transaction);
                if (error != B_OK)
                        RETURN_ERROR(error);

                fFreeBlocks += toFree;
                remaining -= toFree;
                base += toFree;
                groupOffset = 0;
        }

        return B_OK;
}


status_t
BlockAllocator::_FreeInGroup(uint64 base, uint32 count,
        Transaction& transaction)
{
        if (count == 0)
                return B_OK;

        PRINT("BlockAllocator::_FreeInGroup(%" B_PRIu64 ", %" B_PRIu32 ")\n",
                base, count);

        ASSERT(count <= kBlocksPerGroup);
        ASSERT(base % kBlocksPerGroup + count <= kBlocksPerGroup);

        Block block;
        if (!block.GetWritable(fVolume,
                        fAllocationGroupBlock + base / kBlocksPerGroup, transaction)) {
                RETURN_ERROR(B_ERROR);
        }

        uint16* counts = (uint16*)block.Data();

        uint32 blockIndex = base / kBlocksPerBitmapBlock % kBitmapBlocksPerGroup;
        uint64 inBlockOffset = base % kBlocksPerBitmapBlock;
        uint64 remaining = count;

        while (remaining > 0) {
                uint32 toFree = std::min(remaining,
                        kBlocksPerBitmapBlock - inBlockOffset);

                if (counts[blockIndex] + toFree > kBlocksPerBitmapBlock)
                        RETURN_ERROR(B_BAD_VALUE);

                status_t error = _FreeInBitmapBlock(base, toFree, transaction);
                if (error != B_OK)
                        RETURN_ERROR(error);

                counts[blockIndex] += toFree;
                remaining -= toFree;
                base += toFree;
                blockIndex++;
                inBlockOffset = 0;
        }

        return B_OK;
}


status_t
BlockAllocator::_FreeInBitmapBlock(uint64 base, uint32 count,
        Transaction& transaction)
{
        PRINT("BlockAllocator::_FreeInBitmapBlock(%" B_PRIu64 ", %" B_PRIu32 ")\n",
                base, count);

        ASSERT(count <= kBlocksPerBitmapBlock);
        ASSERT(base % kBlocksPerBitmapBlock + count <= kBlocksPerBitmapBlock);

        Block block;
        if (!block.GetWritable(fVolume,
                        fBitmapBlock + base / kBlocksPerBitmapBlock, transaction)) {
                RETURN_ERROR(B_ERROR);
        }

        uint32* bits = (uint32*)block.Data() + base % kBlocksPerBitmapBlock / 32;
        uint32 bitOffset = base % 32;
        uint32 remaining = count;

        // handle partial uint32 at the beginning
        if (bitOffset > 0) {
                uint32 endOffset = std::min(bitOffset + remaining, (uint32)32);

                uint32 mask = ~(((uint32)1 << bitOffset) - 1);
                if (endOffset < 32)
                        mask &= ((uint32)1 << endOffset) - 1;

                if ((*bits & mask) != mask)
                        RETURN_ERROR(B_BAD_VALUE);

                *bits &= ~mask;
                remaining -= endOffset - bitOffset;
        }

        // handle complete uint32s in the middle
        while (remaining >= 32) {
                if (*bits != 0xffffffff)
                        RETURN_ERROR(B_BUSY);

                *bits = 0;
                remaining -= 32;
        }

        // handle partial uint32 at the end
        if (remaining > 0) {
                uint32 mask = ((uint32)1 << remaining) - 1;

                if ((*bits & mask) != mask)
                        return B_BUSY;

                *bits &= ~mask;
        }

        return B_OK;
}


status_t
BlockAllocator::_UpdateSuperBlock(Transaction& transaction)
{
        // write the superblock
        Block block;
        if (!block.GetWritable(fVolume, kCheckSumFSSuperBlockOffset / B_PAGE_SIZE,
                        transaction)) {
                return B_ERROR;
        }

        SuperBlock* superBlock = (SuperBlock*)block.Data();
        superBlock->SetFreeBlocks(fFreeBlocks);

        block.Put();

        return B_OK;
}