// bitscale -- throwaway spike tool.
//
// Rewrites every non-constant getelementptr so that its byte offset is
// multiplied by 8.  The result is an IR in which every pointer value is a
// *bit* address, while loads/stores still carry their normal LLVM byte-sized
// types.  A real TMS34010 backend would lower "load i8 at bit address A" as
// MOVB [A] and "load i32 at bit address A" as MOVE [A],FS0-in-32-bit-mode.
//
// Build:
//   clang++-19 -std=c++17 $(llvm-config-19 --cxxflags) bitscale.cpp \
//       $(llvm-config-19 --ldflags --libs core irreader support) -o bitscale

#include "llvm/IR/Constants.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/DataLayout.h"
#include "llvm/IR/GetElementPtrTypeIterator.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/Operator.h"
#include "llvm/IRReader/IRReader.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/Support/raw_ostream.h"

#include <map>
#include <set>
#include <string>
#include <vector>

using namespace llvm;

struct Stats {
    unsigned geps = 0;
    unsigned const_geps = 0;
    unsigned vector_geps = 0;
    unsigned allocas = 0;
    unsigned memintrinsics = 0;
    unsigned vaarg = 0;
    unsigned byval_args = 0;
    unsigned indirect_calls = 0;
    unsigned loads = 0;
    unsigned stores = 0;
    unsigned ptr_subs = 0;
    unsigned const_ptr_subs = 0;
};

static void count_constants(Value *V, Stats &S, unsigned depth = 0)
{
    if (!V || depth > 64)
        return;
    auto *C = dyn_cast<Constant>(V);
    if (!C)
        return;
    if (auto *CE = dyn_cast<ConstantExpr>(C)) {
        if (CE->getOpcode() == Instruction::Sub) {
            auto *L = dyn_cast<ConstantExpr>(CE->getOperand(0));
            auto *R = dyn_cast<ConstantExpr>(CE->getOperand(1));
            if (L && R && L->getOpcode() == Instruction::PtrToInt &&
                R->getOpcode() == Instruction::PtrToInt)
                ++S.const_ptr_subs;
        }
    }
    if (auto *GEP = dyn_cast<GEPOperator>(C)) {
        ++S.const_geps;
        count_constants(GEP->getPointerOperand(), S, depth + 1);
        for (auto &Op : GEP->indices())
            count_constants(Op.get(), S, depth + 1);
        return;
    }
    for (auto &Op : C->operands())
        count_constants(Op.get(), S, depth + 1);
}

static void count_constant_geps(const Constant *C, Stats &S, unsigned depth = 0)
{
    count_constants(const_cast<Constant *>(C), S, depth);
}

static Stats collect(Module &M)
{
    Stats S;
    for (Function &F : M) {
        for (Argument &A : F.args())
            if (A.hasAttribute(Attribute::ByVal))
                ++S.byval_args;
        for (BasicBlock &BB : F) {
            for (Instruction &I : BB) {
                if (auto *GEP = dyn_cast<GetElementPtrInst>(&I)) {
                    if (GEP->getType()->isVectorTy())
                        ++S.vector_geps;
                    else
                        ++S.geps;
                }
                if (isa<AllocaInst>(&I))
                    ++S.allocas;
                if (isa<LoadInst>(&I))
                    ++S.loads;
                if (isa<StoreInst>(&I))
                    ++S.stores;
                if (auto *BO = dyn_cast<BinaryOperator>(&I)) {
                    if (BO->getOpcode() == Instruction::Sub &&
                        isa<PtrToIntOperator>(BO->getOperand(0)) &&
                        isa<PtrToIntOperator>(BO->getOperand(1)))
                        ++S.ptr_subs;
                }
                if (auto *CB = dyn_cast<CallBase>(&I)) {
                    if (Function *Callee = CB->getCalledFunction()) {
                        if (Callee->isIntrinsic()) {
                            StringRef N = Callee->getName();
                            if (N.starts_with("llvm.memcpy") || N.starts_with("llvm.memmove") ||
                                N.starts_with("llvm.memset"))
                                ++S.memintrinsics;
                            if (N.starts_with("llvm.va_start") || N.starts_with("llvm.va_arg"))
                                ++S.vaarg;
                        }
                    } else {
                        ++S.indirect_calls;
                    }
                    for (unsigned i = 0; i < CB->arg_size(); ++i)
                        if (CB->paramHasAttr(i, Attribute::ByVal))
                            ++S.byval_args;
                }
                for (Value *Op : I.operands())
                    count_constants(Op, S);
            }
        }
    }
    for (GlobalVariable &G : M.globals())
        if (G.hasInitializer())
            count_constant_geps(G.getInitializer(), S);
    return S;
}

static bool scale_constant_gep(ConstantExpr *CE, const DataLayout &DL)
{
    if (CE->getType()->isVectorTy())
        return false;

    auto *GEP = cast<GEPOperator>(CE);
    SmallVector<Value *, 8> Indices;
    for (auto &Op : GEP->indices())
        Indices.push_back(Op.get());
    for (Value *Idx : Indices)
        if (!isa<ConstantInt>(Idx))
            return false; // cannot fold a non-constant constant-expression GEP

    int64_t Off = DL.getIndexedOffsetInType(GEP->getSourceElementType(), Indices);
    Constant *Base = cast<Constant>(GEP->getPointerOperand());
    Constant *BaseInt = ConstantExpr::getPtrToInt(Base, Type::getInt32Ty(CE->getContext()));
    Constant *BitOff = ConstantInt::get(Type::getInt32Ty(CE->getContext()), Off * 8);
    Constant *BitAddr = ConstantExpr::getAdd(BaseInt, BitOff);
    Constant *New = ConstantExpr::getIntToPtr(BitAddr, CE->getType());
    CE->replaceAllUsesWith(New);
    return true;
}

// base_int + (byte offset computed from the GEP's indices) << 3
static bool scale_gep(GetElementPtrInst *GEP)
{
    if (GEP->getType()->isVectorTy())
        return false;
    Value *Base = GEP->getPointerOperand();
    Type *PtrTy = GEP->getType();
    const DataLayout &DL = GEP->getModule()->getDataLayout();

    IRBuilder<> B(GEP);
    Value *BaseInt = B.CreatePtrToInt(Base, B.getInt32Ty(), "bit.base");
    Value *ByteOff = ConstantInt::get(B.getInt32Ty(), 0);

    for (auto GTI = gep_type_begin(GEP), GTE = gep_type_end(GEP); GTI != GTE; ++GTI) {
        Value *Idx = GTI.getOperand();
        Value *Idx32 = B.CreateSExtOrTrunc(Idx, B.getInt32Ty());
        if (StructType *STy = GTI.getStructTypeOrNull()) {
            unsigned FieldNo = cast<ConstantInt>(Idx)->getZExtValue();
            uint64_t FieldOff = DL.getStructLayout(STy)->getElementOffset(FieldNo);
            ByteOff = B.CreateAdd(ByteOff, ConstantInt::get(B.getInt32Ty(), FieldOff));
        } else {
            uint64_t Stride = GTI.getSequentialElementStride(DL);
            Value *Scaled = B.CreateMul(Idx32, ConstantInt::get(B.getInt32Ty(), Stride));
            ByteOff = B.CreateAdd(ByteOff, Scaled);
        }
    }

    Value *BitOff = B.CreateShl(ByteOff, 3, "bit.off");
    Value *BitAddr = B.CreateAdd(BaseInt, BitOff, "bit.addr");
    Value *NewPtr = B.CreateIntToPtr(BitAddr, PtrTy, "bit.ptr");

    GEP->replaceAllUsesWith(NewPtr);
    GEP->eraseFromParent();
    return true;
}

static void collect_const_geps(Value *V, std::vector<ConstantExpr *> &Out,
                               unsigned depth = 0)
{
    if (!V || depth > 64)
        return;
    auto *C = dyn_cast<Constant>(V);
    if (!C)
        return;
    if (auto *CE = dyn_cast<ConstantExpr>(C)) {
        if (CE->getOpcode() == Instruction::GetElementPtr && !CE->getType()->isVectorTy()) {
            Out.push_back(CE);
            return;
        }
    }
    for (auto &Op : C->operands())
        collect_const_geps(Op.get(), Out, depth + 1);
}

int main(int argc, char **argv)
{
    if (argc != 3) {
        errs() << "usage: " << argv[0] << " input.ll output.ll\n";
        return 2;
    }

    LLVMContext Ctx;
    SMDiagnostic Err;
    std::unique_ptr<Module> M = parseIRFile(argv[1], Err, Ctx);
    if (!M) {
        Err.print(argv[0], errs());
        return 1;
    }

    Stats Before = collect(*M);

    // Constant-expression GEPs (e.g. `"hello" + 2`) cannot be replaced by an
    // instruction, so rewrite them as constant inttoptr(add(ptrtoint, 8*off)).
    unsigned ConstScaled = 0;
    const DataLayout &DL = M->getDataLayout();
    bool Changed = true;
    while (Changed) {
        Changed = false;
        std::vector<ConstantExpr *> Work;
        for (GlobalVariable &G : M->globals())
            if (G.hasInitializer())
                collect_const_geps(G.getInitializer(), Work);
        for (Function &F : *M)
            for (BasicBlock &BB : F)
                for (Instruction &I : BB)
                    for (Value *Op : I.operands())
                        collect_const_geps(Op, Work);

        std::set<ConstantExpr *> Seen;
        for (ConstantExpr *CE : Work) {
            if (!Seen.insert(CE).second)
                continue;
            if (scale_constant_gep(CE, DL)) {
                ++ConstScaled;
                Changed = true;
            }
        }
    }

    unsigned Scaled = 0;
    for (Function &F : *M) {
        std::vector<GetElementPtrInst *> Work;
        for (BasicBlock &BB : F)
            for (Instruction &I : BB)
                if (auto *GEP = dyn_cast<GetElementPtrInst>(&I))
                    Work.push_back(GEP);
        for (GetElementPtrInst *GEP : Work)
            if (scale_gep(GEP))
                ++Scaled;
    }
    Stats After = collect(*M);

    errs() << "bitscale: scaled " << Scaled << " GEPs, " << ConstScaled
           << " constant GEPs\n";
    errs() << "  before: geps=" << Before.geps << " const_geps=" << Before.const_geps
           << " vector_geps=" << Before.vector_geps << " allocas=" << Before.allocas
           << " memintrinsics=" << Before.memintrinsics << " va=" << Before.vaarg
           << " byval=" << Before.byval_args << " indirect=" << Before.indirect_calls
           << " loads=" << Before.loads << " stores=" << Before.stores
           << " ptr_subs=" << Before.ptr_subs << " const_ptr_subs=" << Before.const_ptr_subs << "\n";
    errs() << "  after:  geps=" << After.geps << " const_geps=" << After.const_geps
           << " vector_geps=" << After.vector_geps << " allocas=" << After.allocas
           << " memintrinsics=" << After.memintrinsics << " va=" << After.vaarg
           << " byval=" << After.byval_args << " indirect=" << After.indirect_calls
           << " loads=" << After.loads << " stores=" << After.stores
           << " ptr_subs=" << After.ptr_subs << " const_ptr_subs=" << After.const_ptr_subs << "\n";

    std::error_code EC;
    raw_fd_ostream Out(argv[2], EC);
    if (EC) {
        errs() << "cannot write " << argv[2] << ": " << EC.message() << "\n";
        return 1;
    }
    M->print(Out, nullptr);
    return 0;
}