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)
            # Accept PHI or non-PHI values equally for returns; by this point
            # PHIs for the current block should have been materialized at the top.
            if tmp0 is not None:
                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:
                # Resolve direct value; PHIは finalize_phis に一任
                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)

        # If still zero-like and we have predecessor snapshots, synthesize a minimal PHI at block head.
        try:
            zero_like = isinstance(ret_val, ir.Constant)
            if zero_like and preds is not None and block_end_values is not None and bb_map is not None and isinstance(value_id, int):
                # Derive current block id from name like 'bb3'
                cur_bid = None
                try:
                    cur_bid = int(str(builder.block.name).replace('bb',''))
                except Exception:
                    cur_bid = None
                if cur_bid is not None:
                    incoming = []
                    for p in preds.get(cur_bid, []):
                        # Skip self-loop
                        if p == cur_bid: continue
                        v = None
                        try:
                            v = block_end_values.get(p, {}).get(value_id)
                        except Exception:
                            v = None
                        if v is None:
                            v = ir.Constant(return_type, 0)
                        bblk = bb_map.get(p)
                        if bblk is not None:
                            incoming.append((v, bblk))
                    if incoming:
                        phi = builder.phi(return_type, name=f"ret_phi_{value_id}")
                        for (v, bblk) in incoming:
                            phi.add_incoming(v, bblk)
                        ret_val = phi
        except Exception:
            pass
        
        # 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)