refactor: centralize free-threaded layout branching

Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>

Author: MilesCranmerBot

src/C/extras.jl

@@ -1,13 +1,57 @@
 asptr(x) = Base.unsafe_convert(PyPtr, x)
 
-Py_Type(x) = Base.GC.@preserve x begin
-    if CTX.is_free_threaded
-        PyPtr(UnsafePtr{PyObjectFT}(asptr(x)).type[!])
+# Free-threaded CPython builds ("3.14t") currently have different C struct layouts,
+# but there is no stable ABI yet. To keep the code manageable, we centralize the
+# branching in a single macro that rewrites type names in the expression.
+const _FT_TYPE_REPLACEMENTS = Dict{Symbol,Symbol}(
+    :PyObject => :PyObjectFT,
+    :PyVarObject => :PyVarObjectFT,
+    :PyMemoryViewObject => :PyMemoryViewObjectFT,
+    :PySimpleObject => :PySimpleObjectFT,
+    # Used from JlWrap/C.jl via `C.@ft`.
+    :PyJuliaValueObject => :PyJuliaValueObjectFT,
+)
+
+_ft_replace(sym::Symbol) = get(_FT_TYPE_REPLACEMENTS, sym, sym)
+
+function _ft_transform(ex)
+    if ex isa Symbol
+        return _ft_replace(ex)
+    elseif ex isa QuoteNode
+        v = ex.value
+        return v isa Symbol ? QuoteNode(_ft_replace(v)) : ex
+    elseif ex isa Expr
+        # Handle dotted refs like `C.PyObject` (Expr(:., ...)).
+        if ex.head === :. && length(ex.args) == 2 && ex.args[2] isa QuoteNode && ex.args[2].value isa Symbol
+            return Expr(:., _ft_transform(ex.args[1]), QuoteNode(_ft_replace(ex.args[2].value)))
+        end
+        return Expr(ex.head, map(_ft_transform, ex.args)...)
     else
-        PyPtr(UnsafePtr{PyObject}(asptr(x)).type[!])
+        return ex
     end
 end
 
+"""
+    @ft expr
+
+Evaluate `expr`, but when `CTX.is_free_threaded` is true (CPython "free-threaded"
+builds), rewrite internal type names like `PyObject` → `PyObjectFT` inside the
+expression.
+
+This keeps free-threaded branching centralized, so we don't scatter `if
+CTX.is_free_threaded` throughout the code.
+"""
+macro ft(ex)
+    ex_ft = _ft_transform(ex)
+    return :(if CTX.is_free_threaded
+        $(esc(ex_ft))
+    else
+        $(esc(ex))
+    end)
+end
+
+Py_Type(x) = Base.GC.@preserve x @ft PyPtr(UnsafePtr{PyObject}(asptr(x)).type[!])
+
 PyObject_Type(x) = Base.GC.@preserve x (t = Py_Type(asptr(x)); Py_IncRef(t); t)
 
 Py_TypeCheck(o, t) = Base.GC.@preserve o t PyType_IsSubtype(Py_Type(asptr(o)), asptr(t))
@@ -16,13 +60,7 @@ Py_TypeCheckFast(o, f::Integer) = Base.GC.@preserve o PyType_IsSubtypeFast(Py_Ty
 PyType_IsSubtypeFast(t, f::Integer) =
     Base.GC.@preserve t Cint(!iszero(PyType_GetFlags(asptr(t)) & f))
 
-PyMemoryView_GET_BUFFER(m) = Base.GC.@preserve m begin
-    if CTX.is_free_threaded
-        Ptr{Py_buffer}(UnsafePtr{PyMemoryViewObjectFT}(asptr(m)).view)
-    else
-        Ptr{Py_buffer}(UnsafePtr{PyMemoryViewObject}(asptr(m)).view)
-    end
-end
+PyMemoryView_GET_BUFFER(m) = Base.GC.@preserve m @ft Ptr{Py_buffer}(UnsafePtr{PyMemoryViewObject}(asptr(m)).view)
 
 PyType_CheckBuffer(t) = Base.GC.@preserve t begin
     getbuf = PyType_GetSlot(asptr(t), Py_bf_getbuffer)
@@ -79,13 +117,7 @@ function PyOS_RunInputHook()
 end
 
 function PySimpleObject_GetValue(::Type{T}, o) where {T}
-    Base.GC.@preserve o begin
-        if CTX.is_free_threaded
-            UnsafePtr{PySimpleObjectFT{T}}(asptr(o)).value[!]
-        else
-            UnsafePtr{PySimpleObject{T}}(asptr(o)).value[!]
-        end
-    end
+    Base.GC.@preserve o @ft UnsafePtr{PySimpleObject{T}}(asptr(o)).value[!]
 end
 
 # FAST REFCOUNTING

src/JlWrap/C.jl

@@ -26,35 +26,23 @@ const PYJLVALUES = []
 # unused indices in PYJLVALUES
 const PYJLFREEVALUES = Int[]
 
-# Choose the layout based on whether Python is free-threaded.
-pyjl_obj_type() = C.CTX.is_free_threaded ? PyJuliaValueObjectFT : PyJuliaValueObject
-
-function pyjl_obj_ptr(o::C.PyPtr)
-    if C.CTX.is_free_threaded
-        return UnsafePtr{PyJuliaValueObjectFT}(o)
-    else
-        return UnsafePtr{PyJuliaValueObject}(o)
-    end
-end
-
 function _pyjl_new(t::C.PyPtr, ::C.PyPtr, ::C.PyPtr)
     alloc = C.PyType_GetSlot(t, C.Py_tp_alloc)
     alloc == C_NULL && return C.PyNULL
     o = ccall(alloc, C.PyPtr, (C.PyPtr, C.Py_ssize_t), t, 0)
     o == C.PyNULL && return C.PyNULL
-    p = pyjl_obj_ptr(o)
-    p.weaklist[] = C.PyNULL
-    p.value[] = 0
+    C.@ft UnsafePtr{PyJuliaValueObject}(o).weaklist[] = C.PyNULL
+    C.@ft UnsafePtr{PyJuliaValueObject}(o).value[] = 0
     return o
 end
 
 function _pyjl_dealloc(o::C.PyPtr)
-    idx = pyjl_obj_ptr(o).value[]
+    idx = C.@ft UnsafePtr{PyJuliaValueObject}(o).value[]
     if idx != 0
         PYJLVALUES[idx] = nothing
         push!(PYJLFREEVALUES, idx)
     end
-    pyjl_obj_ptr(o).weaklist[!] == C.PyNULL || C.PyObject_ClearWeakRefs(o)
+    (C.@ft UnsafePtr{PyJuliaValueObject}(o).weaklist[!]) == C.PyNULL || C.PyObject_ClearWeakRefs(o)
     freeptr = C.PyType_GetSlot(C.Py_Type(o), C.Py_tp_free)
     freeptr == C_NULL || ccall(freeptr, Cvoid, (C.PyPtr,), o)
     nothing
@@ -335,13 +323,12 @@ function init_c()
 
     # Create members for weakref support
     empty!(_pyjlbase_members)
-    objT = pyjl_obj_type()
     push!(
         _pyjlbase_members,
         C.PyMemberDef(
             name = pointer(_pyjlbase_weaklistoffset_name),
             typ = C.Py_T_PYSSIZET,
-            offset = fieldoffset(objT, 3),
+            offset = (C.@ft fieldoffset(PyJuliaValueObject, 3)),
             flags = C.Py_READONLY,
         ),
         C.PyMemberDef(), # NULL terminator
@@ -363,7 +350,7 @@ function init_c()
     # Create PyType_Spec
     _pyjlbase_spec[] = C.PyType_Spec(
         name = pointer(_pyjlbase_name),
-        basicsize = sizeof(objT),
+        basicsize = (C.@ft sizeof(PyJuliaValueObject)),
         flags = C.Py_TPFLAGS_BASETYPE | C.Py_TPFLAGS_HAVE_VERSION_TAG,
         slots = pointer(_pyjlbase_slots),
     )
@@ -380,14 +367,13 @@ function __init__()
     init_c()
 end
 
-PyJuliaValue_IsNull(o) = Base.GC.@preserve o pyjl_obj_ptr(C.asptr(o)).value[] == 0
+PyJuliaValue_IsNull(o) = Base.GC.@preserve o (C.@ft UnsafePtr{PyJuliaValueObject}(C.asptr(o)).value[]) == 0
 
-PyJuliaValue_GetValue(o) = Base.GC.@preserve o PYJLVALUES[pyjl_obj_ptr(C.asptr(o)).value[]]
+PyJuliaValue_GetValue(o) = Base.GC.@preserve o PYJLVALUES[(C.@ft UnsafePtr{PyJuliaValueObject}(C.asptr(o)).value[])]
 
 PyJuliaValue_SetValue(_o, @nospecialize(v)) = Base.GC.@preserve _o begin
     o = C.asptr(_o)
-    p = pyjl_obj_ptr(o)
-    idx = p.value[]
+    idx = C.@ft UnsafePtr{PyJuliaValueObject}(o).value[]
     if idx == 0
         if isempty(PYJLFREEVALUES)
             push!(PYJLVALUES, v)
@@ -396,7 +382,7 @@ PyJuliaValue_SetValue(_o, @nospecialize(v)) = Base.GC.@preserve _o begin
             idx = pop!(PYJLFREEVALUES)
             PYJLVALUES[idx] = v
         end
-        p.value[] = idx
+        C.@ft UnsafePtr{PyJuliaValueObject}(o).value[] = idx
     else
         PYJLVALUES[idx] = v
     end