Skip to content

Commit b4e6399

Browse files
msimacekmsimacek
committed
Store CodeUnits directly in frozen modules
1 parent cee059b commit b4e6399

File tree

4 files changed

+103
-92
lines changed

4 files changed

+103
-92
lines changed

‎graalpython/com.oracle.graal.python.frozen/freeze_modules.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -553,7 +553,8 @@ def lower_camel_case(str):
553553
def freeze_module(src):
554554
with open(src.pyfile, "r", encoding="utf-8") as src_file, open(src.frozenfile, "wb") as binary_file:
555555
code_obj = compile(src_file.read(), f"<frozen {src.id}>", "exec")
556-
marshal.dump(code_obj, binary_file)
556+
# GraalPy note: we don't use marshal here, we just dump the co_code which implicitly marshals the code unit
557+
binary_file.write(code_obj.co_code)
557558

558559

559560
def write_frozen_modules_map(out_file, modules):

‎graalpython/com.oracle.graal.python/src/com/oracle/graal/python/PythonLanguage.java‎

Lines changed: 41 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -569,43 +569,7 @@ protected CallTarget parse(ParsingRequest request) {
569569
if (MIME_TYPE_BYTECODE.equals(source.getMimeType())) {
570570
byte[] bytes = source.getBytes().toByteArray();
571571
CodeUnit code = MarshalModuleBuiltins.deserializeCodeUnit(null, context, bytes);
572-
boolean internal = shouldMarkSourceInternal(context);
573-
// The original file path should be passed as the name
574-
String name = source.getName();
575-
if (name != null && !name.isEmpty()) {
576-
Source textSource = tryLoadSource(context, code, internal, name);
577-
if (textSource == null) {
578-
if (name.startsWith(FROZEN_FILENAME_PREFIX) && name.endsWith(FROZEN_FILENAME_SUFFIX)) {
579-
String id = name.substring(FROZEN_FILENAME_PREFIX.length(), name.length() - FROZEN_FILENAME_SUFFIX.length());
580-
String fs = context.getEnv().getFileNameSeparator();
581-
String path = context.getStdlibHome() + fs + id.replace(".", fs) + J_PY_EXTENSION;
582-
textSource = tryLoadSource(context, code, internal, path);
583-
if (textSource == null) {
584-
path = context.getStdlibHome() + fs + id.replace(".", fs) + fs + "__init__.py";
585-
textSource = tryLoadSource(context, code, internal, path);
586-
}
587-
}
588-
}
589-
if (textSource != null) {
590-
source = textSource;
591-
}
592-
}
593-
if (internal && !source.isInternal()) {
594-
source = Source.newBuilder(source).internal(true).build();
595-
}
596-
RootNode rootNode = null;
597-
598-
if (PythonOptions.ENABLE_BYTECODE_DSL_INTERPRETER) {
599-
if (source.hasBytes()) {
600-
// Force a character-based source so that source sections work as expected.
601-
source = Source.newBuilder(source).content(Source.CONTENT_NONE).build();
602-
}
603-
rootNode = ((BytecodeDSLCodeUnit) code).createRootNode(context, source);
604-
} else {
605-
rootNode = PBytecodeRootNode.create(this, (BytecodeCodeUnit) code, source);
606-
}
607-
608-
return PythonUtils.getOrCreateCallTarget(rootNode);
572+
return callTargetFromBytecode(context, source, code);
609573
}
610574

611575
String mime = source.getMimeType();
@@ -631,6 +595,46 @@ protected CallTarget parse(ParsingRequest request) {
631595
return parse(context, source, type, false, optimize, false, null, FutureFeature.fromFlags(flags));
632596
}
633597

598+
public RootCallTarget callTargetFromBytecode(PythonContext context, Source source, CodeUnit code) {
599+
boolean internal = shouldMarkSourceInternal(context);
600+
// The original file path should be passed as the name
601+
String name = source.getName();
602+
if (name != null && !name.isEmpty()) {
603+
Source textSource = tryLoadSource(context, code, internal, name);
604+
if (textSource == null) {
605+
if (name.startsWith(FROZEN_FILENAME_PREFIX) && name.endsWith(FROZEN_FILENAME_SUFFIX)) {
606+
String id = name.substring(FROZEN_FILENAME_PREFIX.length(), name.length() - FROZEN_FILENAME_SUFFIX.length());
607+
String fs = context.getEnv().getFileNameSeparator();
608+
String path = context.getStdlibHome() + fs + id.replace(".", fs) + J_PY_EXTENSION;
609+
textSource = tryLoadSource(context, code, internal, path);
610+
if (textSource == null) {
611+
path = context.getStdlibHome() + fs + id.replace(".", fs) + fs + "__init__.py";
612+
textSource = tryLoadSource(context, code, internal, path);
613+
}
614+
}
615+
}
616+
if (textSource != null) {
617+
source = textSource;
618+
}
619+
}
620+
if (internal && !source.isInternal()) {
621+
source = Source.newBuilder(source).internal(true).build();
622+
}
623+
RootNode rootNode;
624+
625+
if (PythonOptions.ENABLE_BYTECODE_DSL_INTERPRETER) {
626+
if (source.hasBytes()) {
627+
// Force a character-based source so that source sections work as expected.
628+
source = Source.newBuilder(source).content(Source.CONTENT_NONE).build();
629+
}
630+
rootNode = ((BytecodeDSLCodeUnit) code).createRootNode(context, source);
631+
} else {
632+
rootNode = PBytecodeRootNode.create(this, (BytecodeCodeUnit) code, source);
633+
}
634+
635+
return PythonUtils.getOrCreateCallTarget(rootNode);
636+
}
637+
634638
private static Source tryLoadSource(PythonContext context, CodeUnit code, boolean internal, String path) {
635639
try {
636640
return Source.newBuilder(PythonLanguage.ID, context.getEnv().getPublicTruffleFile(path)).name(code.name.toJavaStringUncached()).internal(internal).build();

‎graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/ImpModuleBuiltins.java‎

Lines changed: 38 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@
9090
import com.oracle.graal.python.builtins.objects.module.PythonModule;
9191
import com.oracle.graal.python.builtins.objects.object.PythonObject;
9292
import com.oracle.graal.python.builtins.objects.str.PString;
93+
import com.oracle.graal.python.compiler.CodeUnit;
9394
import com.oracle.graal.python.compiler.Compiler;
9495
import com.oracle.graal.python.lib.PyObjectGetAttr;
9596
import com.oracle.graal.python.lib.PyObjectLookupAttr;
@@ -131,6 +132,7 @@
131132
import com.oracle.truffle.api.memory.ByteArraySupport;
132133
import com.oracle.truffle.api.nodes.Node;
133134
import com.oracle.truffle.api.profiles.InlinedConditionProfile;
135+
import com.oracle.truffle.api.source.Source;
134136
import com.oracle.truffle.api.strings.TruffleString;
135137
import com.oracle.truffle.api.utilities.TriState;
136138

@@ -157,20 +159,14 @@ private static class FrozenResult {
157159

158160
private static class FrozenInfo {
159161
@SuppressWarnings("unused") final TruffleString name;
160-
final byte[] data;
161-
final int size;
162+
final CodeUnit code;
162163
final boolean isPackage;
163164
final TruffleString origName;
164165
@SuppressWarnings("unused") final boolean isAlias;
165166

166-
FrozenInfo(byte[] data, int size) {
167-
this(null, data, size, false, null, false);
168-
}
169-
170-
FrozenInfo(TruffleString name, byte[] data, int size, boolean isPackage, TruffleString origName, boolean isAlias) {
167+
FrozenInfo(TruffleString name, CodeUnit code, boolean isPackage, TruffleString origName, boolean isAlias) {
171168
this.name = name;
172-
this.data = data;
173-
this.size = size;
169+
this.code = code;
174170
this.isPackage = isPackage;
175171
this.origName = origName;
176172
this.isAlias = isAlias;
@@ -548,39 +544,45 @@ static Object run(VirtualFrame frame, TruffleString name, Object dataObj,
548544
@Cached PRaiseNode raiseNode) {
549545
FrozenInfo info;
550546
if (dataObj != PNone.NONE) {
547+
byte[] bytes;
548+
int size;
551549
try {
552-
info = new FrozenInfo(bufferLib.getInternalOrCopiedByteArray(dataObj), bufferLib.getBufferLength(dataObj));
550+
bytes = bufferLib.getInternalOrCopiedByteArray(dataObj);
551+
size = bufferLib.getBufferLength(dataObj);
553552
} finally {
554553
bufferLib.release(dataObj, frame, indirectCallData);
555554
}
556-
if (info.size == 0) {
555+
if (size == 0) {
557556
/* Does not contain executable code. */
558557
raiseFrozenError(frame, FROZEN_INVALID, name, constructAndRaiseNode.get(inliningTarget));
559558
}
559+
560+
Object code = null;
561+
562+
try {
563+
code = MarshalModuleBuiltins.Marshal.load(context, bytes, size);
564+
} catch (MarshalError | NumberFormatException e) {
565+
raiseFrozenError(frame, FROZEN_INVALID, name, constructAndRaiseNode.get(inliningTarget));
566+
}
567+
568+
if (!isCodeObjectProfile.profile(inliningTarget, code instanceof PCode)) {
569+
throw raiseNode.raise(inliningTarget, TypeError, ErrorMessages.NOT_A_CODE_OBJECT, name);
570+
}
571+
572+
return code;
560573
} else {
561574
FrozenResult result = findFrozen(context, name, equalNode);
562575
FrozenStatus status = result.status;
563576
info = result.info;
564577
raiseFrozenError(frame, status, name, constructAndRaiseNode.get(inliningTarget));
565-
}
566578

567-
Object code = null;
568-
569-
try {
570-
code = MarshalModuleBuiltins.Marshal.load(context, info.data, info.size);
571-
} catch (MarshalError | NumberFormatException e) {
572-
raiseFrozenError(frame, FROZEN_INVALID, name, constructAndRaiseNode.get(inliningTarget));
573-
}
574-
575-
if (!isCodeObjectProfile.profile(inliningTarget, code instanceof PCode)) {
576-
throw raiseNode.raise(inliningTarget, TypeError, ErrorMessages.NOT_A_CODE_OBJECT, name);
579+
RootCallTarget callTarget = createCallTarget(context, info);
580+
return PFactory.createCode(context.getLanguage(), callTarget);
577581
}
578-
579-
return code;
580582
}
581583
}
582584

583-
@Builtin(name = "find_frozen", parameterNames = {"name", "withData"}, minNumOfPositionalArgs = 1, doc = "find_frozen($module, name, /, *, withdata=False)\n" +
585+
@Builtin(name = "find_frozen", parameterNames = {"name", "withdata"}, minNumOfPositionalArgs = 1, doc = "find_frozen($module, name, /, *, withdata=False)\n" +
584586
"--\n" +
585587
"\n" +
586588
"Return info about the corresponding frozen module (if there is one) or None.\n" +
@@ -594,7 +596,7 @@ static Object run(VirtualFrame frame, TruffleString name, Object dataObj,
594596
" the module\'s current name)")
595597
@GenerateNodeFactory
596598
@ArgumentClinic(name = "name", conversion = ArgumentClinic.ClinicConversion.TString)
597-
@ArgumentClinic(name = "withData", conversion = ArgumentClinic.ClinicConversion.Boolean, defaultValue = "false", useDefaultForNone = true)
599+
@ArgumentClinic(name = "withdata", conversion = ArgumentClinic.ClinicConversion.Boolean, defaultValue = "false", useDefaultForNone = true)
598600
abstract static class FindFrozen extends PythonBinaryClinicBuiltinNode {
599601

600602
@Override
@@ -625,7 +627,8 @@ static Object run(VirtualFrame frame, TruffleString name, boolean withData,
625627
PMemoryView data = null;
626628

627629
if (withData) {
628-
data = memoryViewNode.execute(frame, PFactory.createBytes(context.getLanguage(inliningTarget), info.data));
630+
byte[] bytes = MarshalModuleBuiltins.serializeCodeUnit(inliningTarget, context, info.code);
631+
data = memoryViewNode.execute(frame, PFactory.createBytes(context.getLanguage(inliningTarget), bytes));
629632
}
630633

631634
Object[] returnValues = new Object[]{
@@ -690,16 +693,14 @@ public static PythonModule importFrozenModuleObject(Python3Core core, TruffleStr
690693
}
691694
}
692695

693-
PCode code = (PCode) MarshalModuleBuiltins.Marshal.load(core.getContext(), info.data, info.size);
694-
696+
RootCallTarget callTarget = createCallTarget(core.getContext(), info);
695697
PythonModule module = globals == null ? PFactory.createPythonModule(core.getLanguage(), name) : globals;
696698

697699
if (info.isPackage) {
698700
/* Set __path__ to the empty list */
699701
WriteAttributeToPythonObjectNode.getUncached().execute(module, T___PATH__, PFactory.createList(core.getLanguage()));
700702
}
701703

702-
RootCallTarget callTarget = code.getRootCallTarget();
703704
CallDispatchers.SimpleIndirectInvokeNode.executeUncached(callTarget, PArguments.withGlobals(module));
704705

705706
Object origName = info.origName == null ? PNone.NONE : info.origName;
@@ -708,6 +709,12 @@ public static PythonModule importFrozenModuleObject(Python3Core core, TruffleStr
708709
return module;
709710
}
710711

712+
private static RootCallTarget createCallTarget(PythonContext context, FrozenInfo info) {
713+
String name = PythonLanguage.FROZEN_FILENAME_PREFIX + info.name + PythonLanguage.FROZEN_FILENAME_SUFFIX;
714+
Source source = Source.newBuilder("python", "", name).content(Source.CONTENT_NONE).build();
715+
return context.getLanguage().callTargetFromBytecode(context, source, info.code);
716+
}
717+
711718
/*
712719
* CPython's version of this accepts any object and casts, but all Python-level callers use
713720
* argument clinic to convert the name first. The only exception is
@@ -727,20 +734,11 @@ private static FrozenResult findFrozen(PythonContext context, TruffleString name
727734
boolean isAlias = module.getOriginalName() == null || !equalNode.execute(name, module.getOriginalName(), TS_ENCODING);
728735
FrozenInfo info = new FrozenInfo(name,
729736
module.getCode(),
730-
module.getSize(),
731737
module.isPackage(),
732738
module.getOriginalName(),
733739
!isAlias);
734740

735-
if (module.getCode() == null) {
736-
/* It is frozen but marked as un-importable. */
737-
return new FrozenResult(FROZEN_EXCLUDED, info);
738-
}
739-
740-
if (module.getCode()[0] == '\0' || module.getSize() == 0) {
741-
/* Does not contain executable code. */
742-
return new FrozenResult(FROZEN_INVALID, info);
743-
}
741+
// CPython checks for invalid/empty modules here, but we don't generate those
744742

745743
return new FrozenResult(FROZEN_OKAY, info);
746744
}

‎graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/module/PythonFrozenModule.java‎

Lines changed: 22 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright (c) 2022, 2024, Oracle and/or its affiliates. All rights reserved.
2+
* Copyright (c) 2022, 2025, Oracle and/or its affiliates. All rights reserved.
33
* DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
44
*
55
* The Universal Permissive License (UPL), Version 1.0
@@ -47,24 +47,30 @@
4747
import java.io.IOException;
4848
import java.io.InputStream;
4949

50+
import org.graalvm.nativeimage.ImageInfo;
51+
52+
import com.oracle.graal.python.builtins.modules.MarshalModuleBuiltins;
53+
import com.oracle.graal.python.compiler.CodeUnit;
5054
import com.oracle.graal.python.runtime.PythonOptions;
5155
import com.oracle.truffle.api.strings.TruffleString;
5256

5357
public final class PythonFrozenModule {
58+
private final String symbol;
5459
private final TruffleString originalName;
55-
private final byte[] code;
5660
private final boolean isPackage;
61+
private CodeUnit code;
5762

58-
private static byte[] getByteCode(String symbol) {
63+
private void initCode() {
5964
try {
6065
InputStream resourceAsStream = PythonFrozenModule.class.getResourceAsStream("Frozen" + symbol + "." + getSuffix());
6166
if (resourceAsStream != null) {
62-
return resourceAsStream.readAllBytes();
67+
byte[] bytes = resourceAsStream.readAllBytes();
68+
// TODO exception handling
69+
code = MarshalModuleBuiltins.deserializeCodeUnit(null, null, bytes);
6370
}
6471
} catch (IOException e) {
6572
// fall-through
6673
}
67-
return null;
6874
}
6975

7076
private static String getSuffix() {
@@ -76,13 +82,16 @@ private static String getSuffix() {
7682
}
7783

7884
public PythonFrozenModule(String symbol, String originalName, boolean isPackage) {
79-
this(toTruffleStringUncached(originalName), getByteCode(symbol), isPackage);
85+
this(symbol, toTruffleStringUncached(originalName), isPackage);
8086
}
8187

82-
private PythonFrozenModule(TruffleString originalName, byte[] code, boolean isPackage) {
88+
private PythonFrozenModule(String symbol, TruffleString originalName, boolean isPackage) {
89+
this.symbol = symbol;
8390
this.originalName = originalName;
84-
this.code = code;
8591
this.isPackage = isPackage;
92+
if (ImageInfo.inImageBuildtimeCode()) {
93+
initCode();
94+
}
8695
}
8796

8897
public PythonFrozenModule asPackage(boolean flag) {
@@ -93,23 +102,22 @@ public PythonFrozenModule asPackage(boolean flag) {
93102
if (isPackage) {
94103
origName = T_LANGLE.concatUncached(originalName, TS_ENCODING, false);
95104
}
96-
return new PythonFrozenModule(origName, code, flag);
105+
return new PythonFrozenModule(symbol, origName, flag);
97106
}
98107
}
99108

100109
public TruffleString getOriginalName() {
101110
return originalName;
102111
}
103112

104-
public byte[] getCode() {
113+
public CodeUnit getCode() {
114+
if (!ImageInfo.inImageCode() && code == null) {
115+
initCode();
116+
}
105117
return code;
106118
}
107119

108120
public boolean isPackage() {
109121
return isPackage;
110122
}
111-
112-
public int getSize() {
113-
return code != null ? code.length : 0;
114-
}
115123
}

0 commit comments

Comments
 (0)