hakorune/src/llvm_py/instructions/ret.py

"""
Return instruction lowering
Handles void and value returns
"""

import llvmlite.ir as ir
from typing import Dict, Optional, Any

def lower_return(
    builder: ir.IRBuilder,
    value_id: Optional[int],
    vmap: Dict[int, ir.Value],
    return_type: ir.Type,
    resolver=None,
    preds=None,
    block_end_values=None,
    bb_map=None,
    ctx: Optional[Any] = None,
) -> None:
    """
    Lower MIR Return instruction
    
    Args:
        builder: Current LLVM IR builder
        value_id: Optional return value ID
        vmap: Value map
        return_type: Expected return type
    """
    # Prefer BuildCtx maps if provided
    if ctx is not None:
        try:
            if getattr(ctx, 'resolver', None) is not None:
                resolver = ctx.resolver
            if getattr(ctx, 'preds', None) is not None and preds is None:
                preds = ctx.preds
            if getattr(ctx, 'block_end_values', None) is not None and block_end_values is None:
                block_end_values = ctx.block_end_values
            if getattr(ctx, 'bb_map', None) is not None and bb_map is None:
                bb_map = ctx.bb_map
        except Exception:
            pass
    if value_id is None:
        # Void return
        builder.ret_void()
    else:
        # Get return value (prefer resolver)
        ret_val = None
        # Fast path: if vmap has a concrete non-PHI value defined in this block, use it directly
        if isinstance(value_id, int):
            tmp0 = vmap.get(value_id)
            try:
                is_phi0 = hasattr(tmp0, 'add_incoming')
            except Exception:
                is_phi0 = False
            if tmp0 is not None and not is_phi0:
                ret_val = tmp0
        if ret_val is None:
            if resolver is not None and preds is not None and block_end_values is not None and bb_map is not None:
                # Multi-pred special case: construct a PHI at block head and wire per-pred values
                # to keep PHIs grouped at top and avoid late synthesis that violates ordering.
                cur_bid = int(str(builder.block.name).replace('bb','')) if hasattr(builder.block, 'name') else -1
                pred_ids = [p for p in preds.get(cur_bid, []) if p != cur_bid] if isinstance(preds, dict) else []
                if isinstance(value_id, int) and len(pred_ids) > 1 and isinstance(return_type, ir.IntType):
                    # Create PHI at block head
                    btop = ir.IRBuilder(builder.block)
                    try:
                        btop.position_at_start(builder.block)
                    except Exception:
                        pass
                    ph = btop.phi(return_type, name=f"res_phi_{value_id}_{cur_bid}")
                    # Wire per-pred end values
                    for pred_bid in pred_ids:
                        pred_bb = bb_map.get(pred_bid) if isinstance(bb_map, dict) else None
                        if pred_bb is None:
                            continue
                        val = resolver._value_at_end_i64(value_id, pred_bid, preds, block_end_values, vmap, bb_map)
                        if not hasattr(val, 'type'):
                            val = ir.Constant(return_type, 0)
                        ph.add_incoming(val, pred_bb)
                    ret_val = ph
                else:
                    # Resolve direct value
                    if isinstance(return_type, ir.PointerType):
                        ret_val = resolver.resolve_ptr(value_id, builder.block, preds, block_end_values, vmap)
                    else:
                        is_stringish = False
                        if hasattr(resolver, 'is_stringish'):
                            try:
                                is_stringish = resolver.is_stringish(int(value_id))
                            except Exception:
                                is_stringish = False
                        if is_stringish and hasattr(resolver, 'string_ptrs') and int(value_id) in getattr(resolver, 'string_ptrs'):
                            p = resolver.string_ptrs[int(value_id)]
                            i8p = ir.IntType(8).as_pointer()
                            i64 = ir.IntType(64)
                            boxer = None
                            for f in builder.module.functions:
                                if f.name == 'nyash.box.from_i8_string':
                                    boxer = f; break
                            if boxer is None:
                                boxer = ir.Function(builder.module, ir.FunctionType(i64, [i8p]), name='nyash.box.from_i8_string')
                            ret_val = builder.call(boxer, [p], name='ret_ptr2h')
                        else:
                            ret_val = resolver.resolve_i64(value_id, builder.block, preds, block_end_values, vmap, bb_map)
                
        if ret_val is None:
            # Default to vmap (non-PHI) if available
            tmp = vmap.get(value_id)
            try:
                is_phi = hasattr(tmp, 'add_incoming')
            except Exception:
                is_phi = False
            if tmp is not None and not is_phi:
                ret_val = tmp
        if not ret_val:
            # Default based on return type
            if isinstance(return_type, ir.IntType):
                ret_val = ir.Constant(return_type, 0)
            elif isinstance(return_type, ir.DoubleType):
                ret_val = ir.Constant(return_type, 0.0)
            else:
                # Pointer type - null
                ret_val = ir.Constant(return_type, None)
        
        # Type adjustment if needed
        if hasattr(ret_val, 'type') and ret_val.type != return_type:
            if isinstance(return_type, ir.IntType) and ret_val.type.is_pointer:
                # ptr to int
                ret_val = builder.ptrtoint(ret_val, return_type, name="ret_p2i")
            elif isinstance(return_type, ir.PointerType) and isinstance(ret_val.type, ir.IntType):
                # int to ptr
                ret_val = builder.inttoptr(ret_val, return_type, name="ret_i2p")
            elif isinstance(return_type, ir.IntType) and isinstance(ret_val.type, ir.IntType):
                # int to int conversion
                if return_type.width < ret_val.type.width:
                    # Truncate
                    ret_val = builder.trunc(ret_val, return_type)
                elif return_type.width > ret_val.type.width:
                    # Zero extend
                    ret_val = builder.zext(ret_val, return_type)
        
        builder.ret(ret_val)