StaticTools

# This is all StaticCompiler-friendly
using StaticTools

function print_args(argc::Int, argv::Ptr{Ptr{UInt8}})
    # c"..." lets you construct statically-sized, stack allocated `StaticString`s
    # We also have m"..." and MallocString if you want the same thing but on the heap
    printf(c"Argument count is %d:\n", argc)
    for i=1:argc
        # iᵗʰ input argument string
        pᵢ = unsafe_load(argv, i) # Get pointer
        strᵢ = MallocString(pᵢ) # Can wrap to get high-level interface
        println(strᵢ)
        # No need to `free` since we didn't allocate this memory
    end
    println(c"That was fun, see you next time!")
    return 0
end

# Compile executable
using StaticCompiler
filepath = compile_executable(print_args, (Int64, Ptr{Ptr{UInt8}}), "./")

using StaticTools
function times_table(argc::Int, argv::Ptr{Ptr{UInt8}})
    argc == 3 || return printf(c"Incorrect number of command-line arguments\n")
    rows = argparse(Int64, argv, 2)            # First command-line argument
    cols = argparse(Int64, argv, 3)            # Second command-line argument

    buf = MallocSlabBuffer()
    @no_escape buf begin
        M = @alloc(Int, rows, cols)
        for i=1:rows
            for j=1:cols
                M[i,j] = i*j
            end
        end
        printf(M)
    end
    free(buf)
end

using StaticCompiler
filepath = compile_executable(times_table, (Int64, Ptr{Ptr{UInt8}}), "./")

using StaticTools
function times_table_malloc(argc::Int, argv::Ptr{Ptr{UInt8}})
    argc == 3 || return printf(c"Incorrect number of command-line arguments\n")
    rows = argparse(Int64, argv, 2)            # First command-line argument
    cols = argparse(Int64, argv, 3)            # Second command-line argument

    M = MallocArray{Int64}(undef, rows, cols)
    @inbounds for i=1:rows
        for j=1:cols
            M[i,j] = i*j
        end
    end
    printf(M)
    free(M)
end

using StaticCompiler
filepath = compile_executable(times_table_malloc, (Int64, Ptr{Ptr{UInt8}}), "./")

julia> function times_table(argc::Int, argv::Ptr{Ptr{UInt8}})
           argc == 3 || return printf(c"Incorrect number of command-line arguments\n")
           rows = argparse(Int64, argv, 2)            # First command-line argument
           cols = argparse(Int64, argv, 3)            # Second command-line argument

           M = MallocArray{Int64}(undef, rows, cols)
           @inbounds for i=1:rows
               for j=1:cols
                   M[i,j] = i*j
               end
           end
           printf(M)
           # M = reinterpret(Int32, M)  # This line is problematic, gets you a linker error at least for some users or versions
           # println(c"\n\nThe same array, reinterpreted as Int32:")
           # printf(M)
           free(M)
       end
times_table (generic function with 1 method)

julia> filepath = compile_executable(times_table, (Int64, Ptr{Ptr{UInt8}}), "./")
"/Users/user/times_table"

shell> ./times_table 3 3
1	2	3
2	4	6
3	6	9


The same array, reinterpreted as Int32:
1	2	3
0	0	0
2	4	6
0	0	0
3	6	9
0	0	0

julia> function rand_matrix(argc::Int, argv::Ptr{Ptr{UInt8}})
          argc == 3 || return printf(stderrp(), c"Incorrect number of command-line arguments\n")
          rows = argparse(Int64, argv, 2)            # First command-line argument
          cols = argparse(Int64, argv, 3)            # Second command-line argument

          rng = static_rng()

          M = MallocArray{Float64}(undef, rows, cols)
          rand!(rng, M)         # rand(rng) to generate a single number instead
          printf(M)
          free(M)
       end
rand_matrix (generic function with 1 method)

julia> compile_executable(rand_matrix, (Int64, Ptr{Ptr{UInt8}}), "./")
"/Users/user/rand_matrix"

shell> ./rand_matrix 5 5
7.890932e-01    7.532989e-01    8.593202e-01    4.790301e-01    6.464508e-01
5.619692e-01    9.800402e-02    8.545220e-02    5.545224e-02    2.966089e-01
7.021460e-01    4.587692e-01    9.316740e-01    8.736913e-01    8.271038e-01
8.098993e-01    5.368138e-01    3.055373e-02    3.972266e-01    8.146640e-01
8.241520e-01    7.532375e-01    2.969434e-01    9.436580e-01    2.819992e-01

shell> hyperfine './rand_matrix 5 5'
Benchmark 1: ./rand_matrix 5 5
  Time (mean ± σ):       2.6 ms ±   0.4 ms    [User: 0.9 ms, System: 0.0 ms]
  Range (min … max):     1.8 ms …   4.2 ms    501 runs

  Warning: Command took less than 5 ms to complete. Results might be inaccurate.

shell> ls -alh rand_matrix
  -rwxr-xr-x  1 user  staff   8.8K May 22 14:02 rand_matrix

using StaticCompiler
using StaticTools
using LoopVectorization

@inline function mul!(C::MallocArray, A::MallocArray, B::MallocArray)
    @turbo for n ∈ indices((C,B), 2), m ∈ indices((C,A), 1)
        Cmn = zero(eltype(C))
        for k ∈ indices((A,B), (2,1))
            Cmn += A[m,k] * B[k,n]
        end
        C[m,n] = Cmn
    end
    return C
end

function loopvec_matrix(argc::Int, argv::Ptr{Ptr{UInt8}})
    argc == 3 || return printf(stderrp(), c"Incorrect number of command-line arguments\n")
    rows = argparse(Int64, argv, 2)            # First command-line argument
    cols = argparse(Int64, argv, 3)            # Second command-line argument

    # LHS
    A = MallocArray{Float64}(undef, rows, cols)
    @turbo for i ∈ axes(A, 1)
        for j ∈ axes(A, 2)
           A[i,j] = i*j
        end
    end

    # RHS
    B = MallocArray{Float64}(undef, cols, rows)
    @turbo for i ∈ axes(B, 1)
        for j ∈ axes(B, 2)
           B[i,j] = i*j
        end
    end

    # # Matrix multiplication
    C = MallocArray{Float64}(undef, cols, cols)
    mul!(C, B, A)

    # Print to stdout
    printf(C)

    # Clean up matrices
    free(A)
    free(B)
    free(C)
end

# Attempt to compile
path = compile_executable(loopvec_matrix, (Int64, Ptr{Ptr{UInt8}}), "./")

using StaticCompiler
using StaticTools
using LoopVectorization
using Base: RefValue

@inline function mul!(C::MallocArray, A::MallocArray, B::MallocArray)
    @turbo for n ∈ indices((C,B), 2), m ∈ indices((C,A), 1)
        Cmn = zero(eltype(C))
        for k ∈ indices((A,B), (2,1))
            Cmn += A[m,k] * B[k,n]
        end
        C[m,n] = Cmn
    end
    return 0
end

# this will let us accept pointers to MallocArrays
mul!(C::Ref,A::Ref,B::Ref) = mul!(C[], A[], B[])

# Note that we have to specify a contrete type for each argument when compiling!
# So not just any MallocArra but in this case specifically MallocArray{Float64,2}
# (AKA MallocMatrix{Float64})
tt = (RefValue{MallocMatrix{Float64}}, RefValue{MallocMatrix{Float64}}, RefValue{MallocMatrix{Float64}})
compile_shlib(mul!, tt, "./", "mul_inplace", filename="libmul")

struct MallocArray{T,N}
    pointer::Ptr{T}
    length::Int
    size::NTuple{N, Int}
end

using Libdl
lib = Libdl.dlopen("./libmul.$(Libdl.dlext)", Libdl.RTLD_LOCAL)
mul_inplace = Libdl.dlsym(lib, "julia_mul_inplace")

A = MallocArray{Float64}(undef, 10, 10); A .= 1
B = MallocArray{Float64}(undef, 10, 10); B .= 1
C = MallocArray{Float64}(undef, 10, 10); C .= 0

ra, rb, rc = Ref(A), Ref(B), Ref(C)
pa, pb, pc = pointer_from_objref(ra), pointer_from_objref(rb), pointer_from_objref(rc)

ccall(mul_inplace, Int, (Ptr{Nothing}, Ptr{Nothing}, Ptr{Nothing}), pc, pa, pb)

Libdl.dlclose(lib)

using StaticTools, StaticCompiler

function dlmul()
    lib = StaticTools.dlopen(c"./libmul.dylib")
    mul_inplace = StaticTools.dlsym(lib, c"julia_mul_inplace")

    A = MallocArray{Float64}(undef, 5, 5); fill!(A, 1)
    B = MallocArray{Float64}(undef, 5, 5); fill!(B, 1)
    C = MallocArray{Float64}(undef, 5, 5); fill!(C, 0)

    ra, rb, rc = Ref(A), Ref(B), Ref(C)
    GC.@preserve ra rb rc begin
        pa, pb, pc = pointer_from_objref(ra), pointer_from_objref(rb), pointer_from_objref(rc)
        @ptrcall mul_inplace(pc::Ptr{Nothing}, pa::Ptr{Nothing}, pb::Ptr{Nothing})::Int
    end
    StaticTools.dlclose(lib)
    printf(C)
end

compile_executable(dlmul, (), "./")

using StaticTools, StaticCompiler

function dlmul()
    A = MallocArray{Float64}(undef, 5, 5); fill!(A, 1)
    B = MallocArray{Float64}(undef, 5, 5); fill!(B, 1)
    C = MallocArray{Float64}(undef, 5, 5); fill!(C, 0)

    ra, rb, rc = Ref(A), Ref(B), Ref(C)
    GC.@preserve ra rb rc begin
        pa, pb, pc = pointer_from_objref(ra), pointer_from_objref(rb), pointer_from_objref(rc)
        @symbolcall julia_mul_inplace(pc::Ptr{Nothing}, pa::Ptr{Nothing}, pb::Ptr{Nothing})::Int
    end
    printf(C)
end

compile_executable(dlmul, (), "./", cflags=`-lmul -L./`)

shell> ./dlmul
5.000000e+00    5.000000e+00    5.000000e+00    5.000000e+00    5.000000e+00
5.000000e+00    5.000000e+00    5.000000e+00    5.000000e+00    5.000000e+00
5.000000e+00    5.000000e+00    5.000000e+00    5.000000e+00    5.000000e+00
5.000000e+00    5.000000e+00    5.000000e+00    5.000000e+00    5.000000e+00
5.000000e+00    5.000000e+00    5.000000e+00    5.000000e+00    5.000000e+00

julia> using StaticCompiler, StaticTools, StaticMPI

julia> function mpihello(argc, argv)
           MPI_Init(argc, argv)

           comm = MPI_COMM_WORLD
           world_size, world_rank = MPI_Comm_size(comm), MPI_Comm_rank(comm)

           printf((c"Hello from ", world_rank, c" of ", world_size, c" processors!\n"))
           MPI_Finalize()
       end
mpihello (generic function with 1 method)

julia> compile_executable(mpihello, (Int, Ptr{Ptr{UInt8}}), "./";
           cflags=`-lmpi -L/opt/local/lib/mpich-mp/`
           # -lmpi instructs compiler to link against libmpi.so / libmpi.dylib
           # -L/opt/local/lib/mpich-mp/ provides path to my local MPICH installation where libmpi can be found
       )

ld: warning: object file (./mpihello.o) was built for newer OSX version (12.0) than being linked (10.13)
"/Users/me/code/StaticTools.jl/mpihello"

shell> mpiexec -np 4 ./mpihello
Hello from 1 of 4 processors!
Hello from 3 of 4 processors!
Hello from 2 of 4 processors!
Hello from 0 of 4 processors!