//Owen Stenson
//LLVM hw2

/* Note about how this works
 * This method relies on programs having a particular structure, and so only works
 *  with a small subset of possible programs, i.e. those which only do basic
 *  manipulations on global variables.
 *      Viable llvm IRs must be composed of one or more subblocks of the form
 *          * Mandatory load instruction (variable to register)
 *          * Optional load instruction (variable to register)
 *          * Mandatory operation (binary or unary)
 *          * Mandatory store operation (register to variable)
 *      The sample inputs are in this form and it makes things easier to reason about.
 *      In an effort to not build up an enormous data structure, these subblocks
 *       are identified one at a time and modified in place.
 */

#include "llvm/Support/raw_ostream.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/InstVisitor.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
#include <vector>


using namespace llvm;

namespace {
    class Subblock {
        public:
            Subblock();
            int insert(Instruction &I);
            void clear();
            void print();
            int invert();
            Instruction *load1;
            Instruction *load2;
            Instruction *op;
            Instruction *store;
    };
    Subblock::Subblock() : load1(NULL), load2(NULL), op(NULL), store(NULL) {}
    void Subblock::clear() {
        load1 = NULL;
        load2 = NULL;
        op = NULL;
        store = NULL;
    }
    void Subblock::print() {
        LoadInst *p1 = dyn_cast<LoadInst>(load1);
        StoreInst *s  = dyn_cast<StoreInst>(store);

        errs() << "\tLoad 1: ('" << p1->getPointerOperand()->getName() << "'):  `";
        p1->print(errs());
        errs() << "`\n";

        if (load2 == NULL) {
            errs() << "\tLoad 2: NULL\n";
        } else {
            LoadInst *p2 = dyn_cast<LoadInst>(load2);
            errs() << "\tLoad 2: ('" << p2->getPointerOperand()->getName() << "'):  `";
            p2->print(errs());
            errs() << "`\n";
        }

        errs() << "\tOp:  `";
        op->print(errs());
        errs() << "`\n";

        errs() << "\tStore: ('" << s->getPointerOperand()->getName() << "')  `";
        store->print(errs());
        errs() << "`\n";
    }
    int Subblock::insert(Instruction &I) {
        //returns -1 on error
        //returns  0 on success
        //returns +1 on completion (after `store`)
        if (LoadInst *load = dyn_cast<LoadInst>(&I)) {
            if(load1 == NULL) {
                load1 = load;
            } else if(load2 == NULL) {
                load2 = load;
            } else {
                errs() << "Found more than 2 loads in a subblock; you're boned\n";
                return -1;
            }
            return 0;
        } else if (BinaryOperator *bo = dyn_cast<BinaryOperator>(&I)) {
            if(op == NULL) {
                op = bo;
                return 0;
            } else {
                errs() << "Found more than 1 operation in a subblock; you're boned\n";
                return -1;
            }
        } else if (StoreInst *s = dyn_cast<StoreInst>(&I)) {
            if(store == NULL) {
                store = s;
                return 1;
            } else {
                errs() << "Found more than 1 store in a subblock; you're boned\n";
                return -1;
            }
        } else {
            errs() << "Encountered an unknown instruction:\n";
            I.print(errs());
            errs() << "\n";
            return -1;
        }
    }
    int Subblock::invert() {
        //TODO: don't try to modify in place to avoid memory issues
        //  Also Instructions contain a bunch of data about their surroundings
        //  Probably the best bet to build up a new block instead
        IRBuilder<> builder(op);
        Value *op0 = op->getOperand(0);
        Value *op1 = op->getOperand(1);
        Instruction *op_ = NULL;
        //use appropriate operation
        switch(op->getOpcode()) {
           //case Instruction::Add: op_=BinaryOperator::Create(llvm::Instruction::Sub, op0, op1); break;
            //case Instruction::Add: op_=builder.CreateSub(op0, op1); break;
            case Instruction::Add: op_=BinaryOperator::Create(llvm::Instruction::Sub, op0, op1); break;
            //case Instruction::Sub: op_=builder.CreateAdd(op0, op1); break;
            //case Instruction::Mul: op_=builder.CreateSDiv(op0, op1); break;
            case Instruction::Sub: op_=BinaryOperator::Create(llvm::Instruction::Add, op0, op1); break;
            case Instruction::Mul: op_=BinaryOperator::Create(llvm::Instruction::SDiv, op0, op1); break;
            case Instruction::UDiv: //treat both as signed
            case Instruction::SDiv: op_=BinaryOperator::Create(llvm::Instruction::Mul, op0, op1); break;
            default: errs() << "Unknown\n"; return -1;
        }
        //std::swap(op, op_);
        if(op_) {
            for (auto& U : op->uses()) {
                errs() << "tryna make a change\n";
                User *user = U.getUser();
                user->setOperand(U.getOperandNo(), op_);
            }
        }
        //op->swapOperands() ?
        return 0;
    }

    void reverse_subblocks(std::vector<Subblock> subblocks) {
        //reverses the subblocks in place
        //TODO: again, this probably can't be done in place
        //  because Intructions try to reference each other or something
        //  Try to construct a basic block from each subblock's instructions instead,
        //   traversing the vec in reverse order
        auto begin = subblocks.begin();
        auto end = --subblocks.end();
        while(true) {
            errs() << "XXX\n";
            //swap contents of `begin` and `end`
            std::swap(begin->load1, end->load1);
            std::swap(begin->load2, end->load2);
            std::swap(begin->op, end->op);
            std::swap(begin->store, end->store);

            ++begin;
            if(begin == end) { break; }
            ++end;
            if(begin == end) { break; }
        }
    }


    struct Hw2 : public FunctionPass {
        static char ID; // Pass identification, replacement for typeid
        Hw2() : FunctionPass(ID) {}

        bool runOnFunction(Function &F) override {
            //only run on the `foo` function
            if(F.getName() != "foo") {
                return false;
            }

            Subblock subblock;
            std::vector<Subblock> subblocks;

            for (auto &B : F) {
                for (auto &I : B) {
                    errs() << "<<"; I.print(errs()); errs() << ">>\n";
                    //cycle through the instructions of this function
                    if (dyn_cast<ReturnInst>(&I)) {
                        break;
                    }
                    int rc = subblock.insert(I);
                    if (rc < 0) {
                        return true;
                    } else if (rc == 1) {
                        errs() << "Found complete subblock\n";
                        errs() << "USED TO BE:\n";
                        subblock.print();
                        subblock.invert();
                        errs() << "BECAME:\n";
                        subblock.print();
                        subblocks.push_back(subblock);
                        subblock.clear();
                    }
                }
            }
            errs() << "FOO\n";
            for(auto i: subblocks) {
                i.print();
            }
            //reverse_subblocks(subblocks);
            errs() << "\n\n\nRESULT:\n";
            for(auto i: subblocks) {
                i.print();
            }
            //for (auto &B : F) {
            //    for (auto &I : B) {
            //        errs() << "<<"; I.print(errs()); errs() << ">>\n";
            //    }
            //}

            return true;
        }
    };
}


char Hw2::ID = 0;
static RegisterPass<Hw2> X("hw2", "HW2: shitty block inverter");