main.py

import array
import ctypes
import math
from multiprocessing import Pool
import os
import time


def square_root_python(numbers):
    # Use multiprocessing to parallelize the square root computation
    with Pool() as pool:
        result = pool.map(math.sqrt, numbers)
    return result


def square_root_ctypes(numbers_array, lib, func_name):
    """Call a shared library via ctypes.

    Uses array.array instead of a Python list to avoid expensive copying.
    The naive approach `c_double_array(*numbers)` unpacks every element,
    which is O(n) in Python and extremely slow for large lists.

    array.array has a buffer interface that ctypes can access directly
    via from_buffer(), giving us zero-copy access to the underlying C array.
    """
    n = len(numbers_array)
    c_double_array = ctypes.c_double * n

    # Get direct pointer to array.array's underlying buffer (zero-copy)
    input_array = c_double_array.from_buffer(numbers_array)

    # Output array still needs allocation, but only once
    output_array = c_double_array()

    # Call the function
    func = getattr(lib, func_name)
    func(input_array, ctypes.c_int(n), output_array)

    return output_array


def main():
    print("Time to do some calculations on a list!")
    numbers = []
    for i in range(10000000):
        numbers.append(float(i))

    # Convert to array.array for ctypes implementations.
    # This is done outside the timing loop since pybind11/PyO3 also accept
    # Python lists directly (they handle conversion internally).
    # Using array.array gives ctypes zero-copy buffer access, making the
    # comparison fair by measuring actual computation time, not Python overhead.
    numbers_array = array.array("d", numbers)

    # Python version
    print("\n=== Python (multiprocessing) ===")
    start = time.time()
    _ = square_root_python(numbers)
    end = time.time()
    print(f"Time: {(end - start) * 1000:.2f}ms")

    # C++ version
    try:
        from cpp_bindings import square_root_cpp

        print("\n=== C++ (pybind11) ===")
        start = time.time()
        _ = square_root_cpp(numbers)
        end = time.time()
        print(f"Time: {(end - start) * 1000:.2f}ms")
    except ImportError as e:
        print(f"\n=== C++ (pybind11) ===")
        print(f"Not available: {e}")
        print("Build with: python setup.py build_ext --inplace")

    # Rust version
    try:
        from rust_bindings import square_root_rust

        print("\n=== Rust (PyO3) ===")
        start = time.time()
        _ = square_root_rust(numbers)
        end = time.time()
        print(f"Time: {(end - start) * 1000:.2f}ms")
    except ImportError as e:
        print(f"\n=== Rust (PyO3) ===")
        print(f"Not available: {e}")
        print("Build with: maturin develop --release")

    # Go version
    try:
        script_dir = os.path.dirname(os.path.abspath(__file__))
        lib_path = os.path.join(script_dir, "go_bindings.so")
        go_lib = ctypes.CDLL(lib_path)
        print("\n=== Go (ctypes) ===")
        start = time.time()
        _ = square_root_ctypes(numbers_array, go_lib, "square_root_go")
        end = time.time()
        print(f"Time: {(end - start) * 1000:.2f}ms")
    except OSError as e:
        print(f"\n=== Go (ctypes) ===")
        print(f"Not available: {e}")
        print("Build with: go build -C go_src -buildmode=c-shared -o ../go_bindings.so")

    # Zig version
    try:
        script_dir = os.path.dirname(os.path.abspath(__file__))
        lib_path = os.path.join(script_dir, "zig_bindings.so")
        zig_lib = ctypes.CDLL(lib_path)
        print("\n=== Zig (ctypes) ===")
        start = time.time()
        _ = square_root_ctypes(numbers_array, zig_lib, "square_root_zig")
        end = time.time()
        print(f"Time: {(end - start) * 1000:.2f}ms")
    except OSError as e:
        print(f"\n=== Zig (ctypes) ===")
        print(f"Not available: {e}")
        print(
            "Build with: zig build-lib -O ReleaseFast zig_src/zig_bindings.zig -dynamic -femit-bin=zig_bindings.so"
        )

    # Odin version
    try:
        script_dir = os.path.dirname(os.path.abspath(__file__))
        lib_path = os.path.join(script_dir, "odin_bindings.so")
        odin_lib = ctypes.CDLL(lib_path)
        print("\n=== Odin (ctypes) ===")
        start = time.time()
        _ = square_root_ctypes(numbers_array, odin_lib, "square_root_odin")
        end = time.time()
        print(f"Time: {(end - start) * 1000:.2f}ms")
    except OSError as e:
        print(f"\n=== Odin (ctypes) ===")
        print(f"Not available: {e}")
        print(
            "Build with: odin build odin_src -build-mode:shared -out:odin_bindings.so -o:speed"
        )


if __name__ == "__main__":
    main()