ast.cpp

#include "ast.hpp"
#include <llvm/IR/Verifier.h>
#include <llvm/IR/Constants.h>
#include <llvm/IR/Value.h>
#include <llvm/IR/Type.h>
#include <llvm/IR/Instructions.h>

ASTBlock* ASTNode::programBlock = nullptr;

void ASTNode::initProgramBlock() {
    programBlock = new ASTBlock();
}

ASTBlock* ASTNode::getProgramBlock() {
    return programBlock;
}

void ASTBlock::codegen(CodeGenContext &ctx) {
    for (auto s: statements) {
        s->codegen(ctx);
    }
}

void ASTExprList::codegen(CodeGenContext &ctx) {
    // Usually, we don't need standalone codegen for lists.
    // The parent node (array init, etc.) handles these.
    // No action required.
}

void ASTDimList::codegen(CodeGenContext &ctx) {
    // Dimensions are processed by the parent node (e.g., ASTArrayDecl).
    // No direct codegen needed.
}

void ASTIndexList::codegen(CodeGenContext &ctx) {
    // Indexes are used by ASTArrayAssign or ASTArrayRef.
    // No direct standalone codegen.
}

void ASTSliceDim::codegen(CodeGenContext &ctx) {
    // Slicing is handled by ASTSliceRef during runtime calls.
}

void ASTSliceSpec::codegen(CodeGenContext &ctx) {
    // Handled by ASTSliceRef.
}

void ASTArrayInit::codegen(CodeGenContext &ctx) {
    // Handled in ASTArrayDecl codegen after allocation.
    // This node just holds initial values.
}

void ASTNestedArrayInit::codegen(CodeGenContext &ctx) {
    // Similar to ASTArrayInit, handled by the declaration logic.
}

void ASTDefaultInit::codegen(CodeGenContext &ctx) {
    // Returns the default value for initialization.
    defaultValue->codegen(ctx);
}

void ASTRangeInit::codegen(CodeGenContext &ctx) {
    // range init handled by ASTArrayDecl codegen.
    // We'll just codegen start and end here if needed.
    start->codegen(ctx);
    end->codegen(ctx);
}

void ASTArrayDecl::codegen(CodeGenContext &ctx) {
    // Allocate array using runtime_alloc.
    int ndims = (int)dims->dims.size();
    llvm::Type *i32Ty = ctx.builder.getInt32Ty();
    std::vector<llvm::Constant*> dimConsts;
    for (auto d: dims->dims)
        dimConsts.push_back(llvm::ConstantInt::get(i32Ty, d));

    llvm::ArrayType *arrTy = llvm::ArrayType::get(ctx.builder.getDoubleTy(), 1);

    // Declare runtime_alloc if needed
    llvm::Function *allocFn = ctx.module->getFunction("runtime_alloc");
    if(!allocFn) {
        llvm::Type* intPtrTy = llvm::PointerType::get(i32Ty,0);
        llvm::FunctionType *fty = llvm::FunctionType::get(ctx.getDoublePtrTy(), {i32Ty, intPtrTy}, false);
        allocFn = llvm::Function::Create(fty, llvm::Function::ExternalLinkage, "runtime_alloc", ctx.module.get());
    }

    llvm::GlobalVariable *dimsArr = new llvm::GlobalVariable(
        *ctx.module, llvm::ArrayType::get(i32Ty, ndims), true,
        llvm::GlobalValue::PrivateLinkage,
        llvm::ConstantArray::get(llvm::ArrayType::get(i32Ty, ndims), dimConsts)
    );

    llvm::Value *zero = llvm::ConstantInt::get(i32Ty, 0);
    llvm::Value *dimsPtr = ctx.builder.CreateGEP(dimsArr, {zero, zero});
    llvm::Value *ndimsVal = llvm::ConstantInt::get(i32Ty, ndims);

    llvm::Value *arrPtr = ctx.builder.CreateCall(allocFn, {ndimsVal, dimsPtr});
    ctx.arrays[name] = {arrPtr, dims->dims};

    // If there's an initializer
    if (initializer) {
        // If it's default(...) init
        if (auto defInit = dynamic_cast<ASTDefaultInit*>(initializer)) {
            defInit->defaultValue->codegen(ctx);
            llvm::Value* defVal = ctx.getCurrentValue();

            // runtime_fill(arrPtr, totalSize, defVal)
            int totalSize = 1;
            for (int d : dims->dims) totalSize *= d;
            llvm::Function *fillFn = ctx.module->getFunction("runtime_fill");
            if(!fillFn) {
                llvm::FunctionType *fty = llvm::FunctionType::get(ctx.builder.getVoidTy(),
                    {ctx.getDoublePtrTy(), ctx.builder.getInt32Ty(), ctx.builder.getDoubleTy()}, false);
                fillFn = llvm::Function::Create(fty, llvm::Function::ExternalLinkage, "runtime_fill", ctx.module.get());
            }
            llvm::Value* tsVal = llvm::ConstantInt::get(ctx.builder.getInt32Ty(), totalSize);
            ctx.builder.CreateCall(fillFn, {arrPtr, tsVal, defVal});
        } else if (auto rInit = dynamic_cast<ASTRangeInit*>(initializer)) {
            // range(start,end)
            rInit->start->codegen(ctx);
            llvm::Value *startVal = ctx.getCurrentValue();
            rInit->end->codegen(ctx);
            llvm::Value *endVal = ctx.getCurrentValue();

            int totalSize=1;
            for (int d : dims->dims) totalSize *= d;

            // runtime_range_fill(arrPtr, totalSize, startVal, endVal)
            llvm::Function *rangeFn = ctx.module->getFunction("runtime_range_fill");
            if(!rangeFn) {
                llvm::FunctionType *fty = llvm::FunctionType::get(ctx.builder.getVoidTy(),
                    {ctx.getDoublePtrTy(), ctx.builder.getInt32Ty(), ctx.builder.getDoubleTy(), ctx.builder.getDoubleTy()}, false);
                rangeFn = llvm::Function::Create(fty, llvm::Function::ExternalLinkage, "runtime_range_fill", ctx.module.get());
            }
            llvm::Value *tsVal = llvm::ConstantInt::get(ctx.builder.getInt32Ty(), totalSize);
            ctx.builder.CreateCall(rangeFn, {arrPtr, tsVal, startVal, endVal});
        } else if (auto arrInit = dynamic_cast<ASTArrayInit*>(initializer)) {
            // A static list of values
            // You would iterate over exprs and store them.
            // For simplicity, just assume 1D and store sequentially:
            int totalSize=1;
            for (int d : dims->dims) totalSize *= d;

            llvm::Value* basePtr = arrPtr;
            int idx=0;
            for (auto e: arrInit->list->exprs) {
                if (idx >= totalSize) break;
                e->codegen(ctx);
                llvm::Value *val = ctx.getCurrentValue();
                // runtime_store_element(arrPtr, idx, val)
                llvm::Function *storeFn = ctx.module->getFunction("runtime_store_element");
                if(!storeFn) {
                    llvm::FunctionType *fty = llvm::FunctionType::get(ctx.builder.getVoidTy(),
                        {ctx.getDoublePtrTy(), ctx.builder.getInt32Ty(), ctx.builder.getDoubleTy()}, false);
                    storeFn = llvm::Function::Create(fty, llvm::Function::ExternalLinkage, "runtime_store_element", ctx.module.get());
                }
                llvm::Value* idxVal = llvm::ConstantInt::get(ctx.builder.getInt32Ty(), idx);
                ctx.builder.CreateCall(storeFn, {arrPtr, idxVal, val});
                idx++;
            }
        } else {
            // Nested arrays etc. Similar logic to above (omitted for brevity).
        }
    }

    ctx.setCurrentValue(arrPtr);
}

void ASTIntLit::codegen(CodeGenContext &ctx) {
    ctx.setCurrentValue(llvm::ConstantInt::get(ctx.builder.getInt32Ty(), value));
}

void ASTDoubleLit::codegen(CodeGenContext &ctx) {
    ctx.setCurrentValue(llvm::ConstantFP::get(ctx.builder.getDoubleTy(), value));
}

void ASTVarRef::codegen(CodeGenContext &ctx) {
    // If this language had scalar variables, we'd load them here.
    // For arrays, we just return the pointer stored in ctx.arrays.
    auto it = ctx.arrays.find(name);
    if (it != ctx.arrays.end()) {
        ctx.setCurrentValue(it->second.alloc);
    } else {
        fprintf(stderr, "Unknown variable: %s\n", name.c_str());
        exit(1);
    }
}

void ASTArrayAssign::codegen(CodeGenContext &ctx) {
    // Assign arr[i][j] = value
    auto it = ctx.arrays.find(name);
    if(it==ctx.arrays.end()){
        fprintf(stderr,"Undefined array %s\n", name.c_str());
        exit(1);
    }

    // Compute flat index for indexing
    int ndims = (int)it->second.dims.size();
    int totalSize=1;
    for (auto d: it->second.dims) totalSize*=d;

    // Compute flat index:
    int offset=0;
    int stride=1;
    // For row-major: offset = sum(indices[k] * product_of_dims_for_lower_dimensions)
    // For simplicity, do it in code:
    int mul = 1;
    for (int i = ndims-1; i>=0; i--) {
        mul = 1;
        for (int j=i+1;j<ndims;j++){
            mul *= it->second.dims[j];
        }
        offset += indices->indices[i]*mul;
    }

    // Evaluate value
    value->codegen(ctx);
    llvm::Value *val = ctx.getCurrentValue();

    llvm::Function *storeFn = ctx.module->getFunction("runtime_store_element");
    if(!storeFn) {
        llvm::FunctionType *fty = llvm::FunctionType::get(ctx.builder.getVoidTy(),
            {ctx.getDoublePtrTy(), ctx.builder.getInt32Ty(), ctx.builder.getDoubleTy()}, false);
        storeFn = llvm::Function::Create(fty, llvm::Function::ExternalLinkage, "runtime_store_element", ctx.module.get());
    }

    llvm::Value *arrPtr = it->second.alloc;
    llvm::Value *idxVal = llvm::ConstantInt::get(ctx.builder.getInt32Ty(), offset);
    ctx.builder.CreateCall(storeFn, {arrPtr, idxVal, val});
}

void ASTResizeCall::codegen(CodeGenContext &ctx) {
    // resize(arr, d1, d2, ...)
    auto it = ctx.arrays.find(arrayName);
    if(it==ctx.arrays.end()){
        fprintf(stderr, "Undefined array: %s\n", arrayName.c_str());
        exit(1);
    }

    int new_ndims = (int)newDims.size();

    llvm::Type *i32Ty = ctx.builder.getInt32Ty();
    std::vector<llvm::Constant*> dimConsts;
    for(auto d: newDims)
        dimConsts.push_back(llvm::ConstantInt::get(i32Ty, d));

    llvm::GlobalVariable *dimsArr = new llvm::GlobalVariable(
        *ctx.module, llvm::ArrayType::get(i32Ty, new_ndims), true,
        llvm::GlobalValue::PrivateLinkage,
        llvm::ConstantArray::get(llvm::ArrayType::get(i32Ty, new_ndims), dimConsts)
    );

    llvm::Value *zero = llvm::ConstantInt::get(i32Ty,0);
    llvm::Value *dimsPtr = ctx.builder.CreateGEP(dimsArr, {zero, zero});
    llvm::Value *ndimsVal = llvm::ConstantInt::get(i32Ty, new_ndims);

    llvm::Function *resizeFn = ctx.module->getFunction("runtime_resize");
    if(!resizeFn) {
        llvm::Type *intPtrTy = llvm::PointerType::get(i32Ty,0);
        llvm::FunctionType *fty = llvm::FunctionType::get(ctx.getDoublePtrTy(),
            {ctx.getDoublePtrTy(), i32Ty, intPtrTy}, false);
        // This signature may differ depending on how you define runtime_resize.
        // Adjust accordingly.
        resizeFn = llvm::Function::Create(fty, llvm::Function::ExternalLinkage, "runtime_resize", ctx.module.get());
    }

    llvm::Value *oldPtr = it->second.alloc;
    // Call runtime_resize(oldPtr, new_ndims, dimsPtr)
    llvm::Value *newPtr = ctx.builder.CreateCall(resizeFn, {oldPtr, ndimsVal, dimsPtr});

    it->second.alloc = newPtr;
    it->second.dims = newDims;
    ctx.setCurrentValue(newPtr);
}

void ASTLengthCall::codegen(CodeGenContext &ctx) {
    auto it = ctx.arrays.find(arrName);
    if(it==ctx.arrays.end()){
        fprintf(stderr,"Undefined array: %s\n", arrName.c_str());
        exit(1);
    }
    // return the size of the given dimension
    if (dim < 1 || dim > (int)it->second.dims.size()) {
        fprintf(stderr, "Dimension out of range in length(%s,%d)\n", arrName.c_str(), dim);
        exit(1);
    }
    int size = it->second.dims[dim-1];
    ctx.setCurrentValue(llvm::ConstantInt::get(ctx.builder.getInt32Ty(), size));
}

void ASTSumCall::codegen(CodeGenContext &ctx) {
    auto it = ctx.arrays.find(arrName);
    if(it==ctx.arrays.end()){
        fprintf(stderr,"Undefined array %s\n", arrName.c_str());
        exit(1);
    }
    int totalSize=1;
    for(auto d: it->second.dims) totalSize*=d;
    llvm::Value *ptr = it->second.alloc;

    llvm::Function *sumFn = ctx.module->getFunction("runtime_sum");
    if(!sumFn) {
        llvm::FunctionType *fty = llvm::FunctionType::get(ctx.builder.getDoubleTy(),
            {ctx.getDoublePtrTy(), ctx.builder.getInt32Ty()}, false);
        sumFn = llvm::Function::Create(fty, llvm::Function::ExternalLinkage, "runtime_sum", ctx.module.get());
    }
    llvm::Value *tsVal = llvm::ConstantInt::get(ctx.builder.getInt32Ty(), totalSize);
    llvm::Value *res = ctx.builder.CreateCall(sumFn, {ptr, tsVal});
    ctx.setCurrentValue(res);
}

void ASTMaxCall::codegen(CodeGenContext &ctx) {
    auto it = ctx.arrays.find(arrName);
    if(it==ctx.arrays.end()){
        fprintf(stderr,"Undefined array %s\n", arrName.c_str());
        exit(1);
    }
    int totalSize=1;
    for(auto d: it->second.dims) totalSize*=d;

    llvm::Function *maxFn = ctx.module->getFunction("runtime_max");
    if(!maxFn) {
        llvm::FunctionType *fty = llvm::FunctionType::get(ctx.builder.getDoubleTy(),
            {ctx.getDoublePtrTy(), ctx.builder.getInt32Ty()}, false);
        maxFn = llvm::Function::Create(fty, llvm::Function::ExternalLinkage, "runtime_max", ctx.module.get());
    }

    llvm::Value *ptr = it->second.alloc;
    llvm::Value *tsVal = llvm::ConstantInt::get(ctx.builder.getInt32Ty(), totalSize);
    llvm::Value *res = ctx.builder.CreateCall(maxFn, {ptr, tsVal});
    ctx.setCurrentValue(res);
}

void ASTAverageCall::codegen(CodeGenContext &ctx) {
    // avg = sum/totalSize
    auto it = ctx.arrays.find(arrName);
    if(it==ctx.arrays.end()){
        fprintf(stderr,"Undefined array %s\n", arrName.c_str());
        exit(1);
    }
    int totalSize=1;
    for(auto d: it->second.dims) totalSize*=d;

    // runtime_average(arrPtr, totalSize)
    llvm::Function *avgFn = ctx.module->getFunction("runtime_average");
    if(!avgFn) {
        llvm::FunctionType *fty = llvm::FunctionType::get(ctx.builder.getDoubleTy(),
            {ctx.getDoublePtrTy(), ctx.builder.getInt32Ty()}, false);
        avgFn = llvm::Function::Create(fty, llvm::Function::ExternalLinkage, "runtime_average", ctx.module.get());
    }

    llvm::Value *ptr = it->second.alloc;
    llvm::Value *tsVal = llvm::ConstantInt::get(ctx.builder.getInt32Ty(), totalSize);
    llvm::Value *res = ctx.builder.CreateCall(avgFn, {ptr, tsVal});
    ctx.setCurrentValue(res);
}

void ASTMapCall::codegen(CodeGenContext &ctx) {
    // map(arr, func)
    auto it = ctx.arrays.find(arrName);
    if(it==ctx.arrays.end()){
        fprintf(stderr,"Undefined array %s\n", arrName.c_str());
        exit(1);
    }

    int totalSize=1;
    for(auto d: it->second.dims) totalSize*=d;

    // runtime_map(arrPtr, totalSize, funcName)
    // Assuming runtime_map(double*, int, char*) -> double*
    llvm::Function *mapFn = ctx.module->getFunction("runtime_map");
    if(!mapFn) {
        llvm::Type *charPtrTy = llvm::Type::getInt8PtrTy(ctx.context);
        llvm::FunctionType *fty = llvm::FunctionType::get(ctx.getDoublePtrTy(),
            {ctx.getDoublePtrTy(), ctx.builder.getInt32Ty(), charPtrTy}, false);
        mapFn = llvm::Function::Create(fty, llvm::Function::ExternalLinkage, "runtime_map", ctx.module.get());
    }

    llvm::Value *ptr = it->second.alloc;
    llvm::Value *tsVal = llvm::ConstantInt::get(ctx.builder.getInt32Ty(), totalSize);
    llvm::Value *funcNameVal = ctx.builder.CreateGlobalStringPtr(funcName);
    llvm::Value *res = ctx.builder.CreateCall(mapFn, {ptr, tsVal, funcNameVal});
    it->second.alloc = res;
    ctx.setCurrentValue(res);
}

void ASTReduceCall::codegen(CodeGenContext &ctx) {
    // reduce(arr, func)
    auto it = ctx.arrays.find(arrName);
    if(it==ctx.arrays.end()){
        fprintf(stderr,"Undefined array %s\n", arrName.c_str());
        exit(1);
    }

    int totalSize=1;
    for(auto d: it->second.dims) totalSize*=d;

    // runtime_reduce(arrPtr, totalSize, funcName) -> double
    llvm::Function *reduceFn = ctx.module->getFunction("runtime_reduce");
    if(!reduceFn) {
        llvm::Type *charPtrTy = llvm::Type::getInt8PtrTy(ctx.context);
        llvm::FunctionType *fty = llvm::FunctionType::get(ctx.builder.getDoubleTy(),
            {ctx.getDoublePtrTy(), ctx.builder.getInt32Ty(), charPtrTy}, false);
        reduceFn = llvm::Function::Create(fty, llvm::Function::ExternalLinkage, "runtime_reduce", ctx.module.get());
    }

    llvm::Value *ptr = it->second.alloc;
    llvm::Value *tsVal = llvm::ConstantInt::get(ctx.builder.getInt32Ty(), totalSize);
    llvm::Value *funcNameVal = ctx.builder.CreateGlobalStringPtr(funcName);
    llvm::Value *res = ctx.builder.CreateCall(reduceFn, {ptr, tsVal, funcNameVal});
    ctx.setCurrentValue(res);
}

void ASTFlattenCall::codegen(CodeGenContext &ctx) {
    // flatten(arr)
    auto it = ctx.arrays.find(arrName);
    if(it==ctx.arrays.end()){
        fprintf(stderr,"Undefined array %s\n", arrName.c_str());
        exit(1);
    }

    int ndims = (int)it->second.dims.size();
    int totalSize=1;
    for(auto d: it->second.dims) totalSize*=d;

    // runtime_flatten(arrPtr, ndims, dims) -> double*
    llvm::Function *flatFn = ctx.module->getFunction("runtime_flatten");
    if(!flatFn) {
        llvm::Type *i32Ty = ctx.builder.getInt32Ty();
        llvm::Type *intPtrTy = llvm::PointerType::get(i32Ty,0);
        llvm::FunctionType *fty = llvm::FunctionType::get(ctx.getDoublePtrTy(),
            {ctx.getDoublePtrTy(), i32Ty, intPtrTy}, false);
        flatFn = llvm::Function::Create(fty, llvm::Function::ExternalLinkage, "runtime_flatten", ctx.module.get());
    }

    // Pass dims
    std::vector<llvm::Constant*> dimConsts;
    for(auto d: it->second.dims)
        dimConsts.push_back(llvm::ConstantInt::get(ctx.builder.getInt32Ty(), d));
    llvm::GlobalVariable *dimsArr = new llvm::GlobalVariable(*ctx.module,
        llvm::ArrayType::get(ctx.builder.getInt32Ty(), ndims), true,
        llvm::GlobalValue::PrivateLinkage,
        llvm::ConstantArray::get(llvm::ArrayType::get(ctx.builder.getInt32Ty(), ndims), dimConsts));

    llvm::Value* zero = llvm::ConstantInt::get(ctx.builder.getInt32Ty(),0);
    llvm::Value* dimsPtr = ctx.builder.CreateGEP(dimsArr, {zero, zero});
    llvm::Value *ndimsVal = llvm::ConstantInt::get(ctx.builder.getInt32Ty(), ndims);

    llvm::Value *res = ctx.builder.CreateCall(flatFn, {it->second.alloc, ndimsVal, dimsPtr});
    // Now arr is 1D
    it->second.alloc = res;
    it->second.dims = { totalSize };
    ctx.setCurrentValue(res);
}

void ASTCopyCall::codegen(CodeGenContext &ctx) {
    // copy(arr)
    auto it = ctx.arrays.find(arrName);
    if (it==ctx.arrays.end()) {
        fprintf(stderr,"Undefined array %s\n", arrName.c_str());
        exit(1);
    }

    int ndims = (int)it->second.dims.size();

    llvm::Function *copyFn = ctx.module->getFunction("runtime_copy");
    if(!copyFn) {
        // runtime_copy(double* arr) -> double*
        llvm::FunctionType *fty = llvm::FunctionType::get(ctx.getDoublePtrTy(),
            {ctx.getDoublePtrTy()}, false);
        copyFn = llvm::Function::Create(fty, llvm::Function::ExternalLinkage, "runtime_copy", ctx.module.get());
    }

    llvm::Value *res = ctx.builder.CreateCall(copyFn, {it->second.alloc});
    // This returns a new array with the same data.
    // The dimension info remains the same.
    // We might not store this result back into ctx.arrays since it's just a returned array.
    // The user might assign it to a new variable. Without a var decl, we just return it.
    ctx.setCurrentValue(res);
}

void ASTCloneStructureCall::codegen(CodeGenContext &ctx) {
    // clone_structure(arr)
    auto it = ctx.arrays.find(arrName);
    if(it==ctx.arrays.end()){
        fprintf(stderr,"Undefined array %s\n", arrName.c_str());
        exit(1);
    }

    llvm::Function *cloneFn = ctx.module->getFunction("runtime_clone_structure");
    if(!cloneFn) {
        llvm::Type *i32Ty = ctx.builder.getInt32Ty();
        llvm::Type *intPtrTy = llvm::PointerType::get(i32Ty,0);
        llvm::FunctionType *fty = llvm::FunctionType::get(ctx.getDoublePtrTy(),
            {i32Ty,intPtrTy}, false);
        cloneFn = llvm::Function::Create(fty, llvm::Function::ExternalLinkage, "runtime_clone_structure", ctx.module.get());
    }

    int ndims = (int)it->second.dims.size();
    std::vector<llvm::Constant*> dimConsts;
    for(auto d: it->second.dims)
        dimConsts.push_back(llvm::ConstantInt::get(ctx.builder.getInt32Ty(), d));
    llvm::GlobalVariable *dimsArr = new llvm::GlobalVariable(*ctx.module,
        llvm::ArrayType::get(ctx.builder.getInt32Ty(), ndims), true,
        llvm::GlobalValue::PrivateLinkage,
        llvm::ConstantArray::get(llvm::ArrayType::get(ctx.builder.getInt32Ty(), ndims), dimConsts));

    llvm::Value* zero = llvm::ConstantInt::get(ctx.builder.getInt32Ty(),0);
    llvm::Value* dimsPtr = ctx.builder.CreateGEP(dimsArr, {zero, zero});
    llvm::Value *ndimsVal = llvm::ConstantInt::get(ctx.builder.getInt32Ty(), ndims);

    llvm::Value *res = ctx.builder.CreateCall(cloneFn, {ndimsVal, dimsPtr});
    ctx.setCurrentValue(res);
}

void ASTReshapeCall::codegen(CodeGenContext &ctx) {
    // reshape(arr, d1, d2, ...)
    auto it = ctx.arrays.find(arrName);
    if(it==ctx.arrays.end()){
        fprintf(stderr,"Undefined array %s\n", arrName.c_str());
        exit(1);
    }

    llvm::Function *reshapeFn = ctx.module->getFunction("runtime_reshape");
    if(!reshapeFn) {
        llvm::Type *i32Ty = ctx.builder.getInt32Ty();
        llvm::Type *intPtrTy = llvm::PointerType::get(i32Ty,0);
        llvm::FunctionType *fty = llvm::FunctionType::get(ctx.getDoublePtrTy(),
            {ctx.getDoublePtrTy(), i32Ty, intPtrTy}, false);
        reshapeFn = llvm::Function::Create(fty, llvm::Function::ExternalLinkage, "runtime_reshape", ctx.module.get());
    }

    int new_ndims = (int)dims.size();
    std::vector<llvm::Constant*> dimConsts;
    for(auto d: dims)
        dimConsts.push_back(llvm::ConstantInt::get(ctx.builder.getInt32Ty(), d));

    llvm::GlobalVariable *dimsArr = new llvm::GlobalVariable(
        *ctx.module, llvm::ArrayType::get(ctx.builder.getInt32Ty(), new_ndims), true,
        llvm::GlobalValue::PrivateLinkage,
        llvm::ConstantArray::get(llvm::ArrayType::get(ctx.builder.getInt32Ty(), new_ndims), dimConsts));

    llvm::Value *zero = llvm::ConstantInt::get(ctx.builder.getInt32Ty(),0);
    llvm::Value *dimsPtr = ctx.builder.CreateGEP(dimsArr, {zero, zero});
    llvm::Value *ndimsVal = llvm::ConstantInt::get(ctx.builder.getInt32Ty(), new_ndims);

    llvm::Value *res = ctx.builder.CreateCall(reshapeFn, {it->second.alloc, ndimsVal, dimsPtr});
    it->second.alloc = res;
    it->second.dims = dims;
    ctx.setCurrentValue(res);
}

void ASTSerializeCall::codegen(CodeGenContext &ctx) {
    // serialize(arr) -> returns a string (char*)
    auto it = ctx.arrays.find(arrName);
    if(it==ctx.arrays.end()){
        fprintf(stderr,"Undefined array %s\n", arrName.c_str());
        exit(1);
    }

    llvm::Function *serFn = ctx.module->getFunction("runtime_serialize");
    if(!serFn) {
        // runtime_serialize(double* arr) -> i8*
        llvm::Type *charPtr = llvm::Type::getInt8PtrTy(ctx.context);
        llvm::FunctionType *fty = llvm::FunctionType::get(charPtr, {ctx.getDoublePtrTy()}, false);
        serFn = llvm::Function::Create(fty, llvm::Function::ExternalLinkage, "runtime_serialize", ctx.module.get());
    }

    llvm::Value *res = ctx.builder.CreateCall(serFn, {it->second.alloc});
    ctx.setCurrentValue(res);
}

void ASTDeserializeCall::codegen(CodeGenContext &ctx) {
    // deserialize(string_expr)
    // string_expr should produce a i8* (C string)
    strExpr->codegen(ctx);
    llvm::Value *strVal = ctx.getCurrentValue();

    llvm::Function *deserFn = ctx.module->getFunction("runtime_deserialize");
    if(!deserFn) {
        // runtime_deserialize(i8*) -> double*
        llvm::Type *charPtr = llvm::Type::getInt8PtrTy(ctx.context);
        llvm::FunctionType *fty = llvm::FunctionType::get(ctx.getDoublePtrTy(), {charPtr}, false);
        deserFn = llvm::Function::Create(fty, llvm::Function::ExternalLinkage, "runtime_deserialize", ctx.module.get());
    }

    llvm::Value *res = ctx.builder.CreateCall(deserFn, {strVal});
    ctx.setCurrentValue(res);
}

void ASTBroadcastOp::codegen(CodeGenContext &ctx) {
    // Broadcasting expr + expr
    // codegen lhs and rhs
    lhs->codegen(ctx);
    llvm::Value *L = ctx.getCurrentValue();
    rhs->codegen(ctx);
    llvm::Value *R = ctx.getCurrentValue();

    // We assume both L and R are arrays. We must get dims from ctx.
    // In a real implementation, we'd find a compatible broadcast shape, etc.

    // For demonstration:
    // runtime_broadcast_add(double* A, int A_ndims, int*A_dims, double* B, int B_ndims, int*B_dims) -> double*
    // Just call a runtime function.

    // For simplicity, assume single dimension for demonstration:
    llvm::Function *broadFn = ctx.module->getFunction("runtime_broadcast_add");
    if(!broadFn) {
        // define a function that takes arrays + metadata
        // We'll just pretend it exists with a simplified signature:
        // runtime_broadcast_add(double* A, double* B) -> double*
        // Adjust as needed.
        llvm::FunctionType *fty = llvm::FunctionType::get(ctx.getDoublePtrTy(),
            {ctx.getDoublePtrTy(), ctx.getDoublePtrTy()}, false);
        broadFn = llvm::Function::Create(fty, llvm::Function::ExternalLinkage, "runtime_broadcast_add", ctx.module.get());
    }

    llvm::Value *res = ctx.builder.CreateCall(broadFn, {L,R});
    ctx.setCurrentValue(res);
}

void ASTSliceRef::codegen(CodeGenContext &ctx) {
    // slicing: arr[...] 
    // runtime_slice(arr, sliceSpec) -> double*
    // This requires complex logic to interpret sliceSpec.
    // We'll just call a runtime function and pass slicing parameters.

    // For simplicity, assume runtime_slice(double* arr, int ndims, int*dims, int sliceCount, slice_struct...) -> double*
    // This is complex; just a placeholder:

    llvm::Function *sliceFn = ctx.module->getFunction("runtime_slice");
    if(!sliceFn) {
        // We'll just define a simple placeholder:
        // runtime_slice(double* arr) -> double*
        llvm::FunctionType *fty = llvm::FunctionType::get(ctx.getDoublePtrTy(), {ctx.getDoublePtrTy()}, false);
        sliceFn = llvm::Function::Create(fty, llvm::Function::ExternalLinkage, "runtime_slice", ctx.module.get());
    }

    auto it = ctx.arrays.find(arrName);
    if(it==ctx.arrays.end()){
        fprintf(stderr,"Undefined array %s\n", arrName.c_str());
        exit(1);
    }

    // Real slicing would require passing sliceSpec info.
    // Just call sliceFn with arrPtr for demonstration:
    llvm::Value *res = ctx.builder.CreateCall(sliceFn, {it->second.alloc});
    ctx.setCurrentValue(res);
}

void ASTExprStmt::codegen(CodeGenContext &ctx) {
    expr->codegen(ctx);
    // Result is not used; it's a statement.
    // Just discard.
}

s