Skip to content

Commit 093f00a

Browse files
jPaolantoniojPaolantonio
authored
[Compile Time Constant Extraction] Extract enums (#63005)
1 parent d951335 commit 093f00a

File tree

4 files changed

+553
-35
lines changed

4 files changed

+553
-35
lines changed

‎include/swift/AST/ConstTypeInfo.h‎

Lines changed: 42 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ class CompileTimeValue {
3434
Dictionary,
3535
Array,
3636
Tuple,
37+
Enum,
3738
Runtime
3839
};
3940

@@ -73,19 +74,19 @@ struct FunctionParameter {
7374
/// with a collection of (potentially compile-time-known) parameters
7475
class InitCallValue : public CompileTimeValue {
7576
public:
76-
InitCallValue(std::string Name, std::vector<FunctionParameter> Parameters)
77-
: CompileTimeValue(ValueKind::InitCall), Name(Name),
78-
Parameters(Parameters) {}
77+
InitCallValue(swift::Type Type, std::vector<FunctionParameter> Parameters)
78+
: CompileTimeValue(ValueKind::InitCall), Type(Type),
79+
Parameters(Parameters) {}
7980

8081
static bool classof(const CompileTimeValue *T) {
8182
return T->getKind() == ValueKind::InitCall;
8283
}
8384

84-
std::string getName() const { return Name; }
85+
swift::Type getType() const { return Type; }
8586
std::vector<FunctionParameter> getParameters() const { return Parameters; }
8687

8788
private:
88-
std::string Name;
89+
swift::Type Type;
8990
std::vector<FunctionParameter> Parameters;
9091
};
9192

@@ -159,6 +160,28 @@ class DictionaryValue : public CompileTimeValue {
159160
std::vector<std::shared_ptr<TupleValue>> Elements;
160161
};
161162

163+
/// An enum value representation
164+
class EnumValue : public CompileTimeValue {
165+
public:
166+
EnumValue(std::string Identifier,
167+
llvm::Optional<std::vector<FunctionParameter>> Parameters)
168+
: CompileTimeValue(ValueKind::Enum), Identifier(Identifier),
169+
Parameters(Parameters) {}
170+
171+
std::string getIdentifier() const { return Identifier; }
172+
llvm::Optional<std::vector<FunctionParameter>> getParameters() const {
173+
return Parameters;
174+
}
175+
176+
static bool classof(const CompileTimeValue *T) {
177+
return T->getKind() == ValueKind::Enum;
178+
}
179+
180+
private:
181+
std::string Identifier;
182+
llvm::Optional<std::vector<FunctionParameter>> Parameters;
183+
};
184+
162185
/// A representation of an arbitrary value that does not fall under
163186
/// any of the above categories.
164187
class RuntimeValue : public CompileTimeValue {
@@ -175,6 +198,19 @@ struct CustomAttrValue {
175198
std::vector<FunctionParameter> Parameters;
176199
};
177200

201+
/// A representation of a single associated value for an enumeration case.
202+
struct EnumElementParameterValue {
203+
llvm::Optional<std::string> Label;
204+
swift::Type Type;
205+
};
206+
207+
/// A representation of a single enumeration case.
208+
struct EnumElementDeclValue {
209+
std::string Name;
210+
llvm::Optional<std::string> RawValue;
211+
llvm::Optional<std::vector<EnumElementParameterValue>> Parameters;
212+
};
213+
178214
struct ConstValueTypePropertyInfo {
179215
swift::VarDecl *VarDecl;
180216
std::shared_ptr<CompileTimeValue> Value;
@@ -194,6 +230,7 @@ struct ConstValueTypePropertyInfo {
194230
struct ConstValueTypeInfo {
195231
swift::NominalTypeDecl *TypeDecl;
196232
std::vector<ConstValueTypePropertyInfo> Properties;
233+
llvm::Optional<std::vector<EnumElementDeclValue>> EnumElements;
197234
};
198235
} // namespace swift
199236
#endif

‎lib/ConstExtract/ConstExtract.cpp‎

Lines changed: 179 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -120,7 +120,29 @@ parseProtocolListFromFile(StringRef protocolListFilePath,
120120
return true;
121121
}
122122

123-
static std::shared_ptr<CompileTimeValue> extractCompileTimeValue(Expr *expr) {
123+
static std::shared_ptr<CompileTimeValue> extractCompileTimeValue(Expr *expr);
124+
125+
static std::vector<FunctionParameter>
126+
extractFunctionArguments(const ArgumentList *args) {
127+
std::vector<FunctionParameter> parameters;
128+
129+
for (auto arg : *args) {
130+
auto argExpr = arg.getExpr();
131+
const auto label = arg.getLabel().str().str();
132+
const auto type = argExpr->getType();
133+
if (auto defaultArgument = dyn_cast<DefaultArgumentExpr>(argExpr)) {
134+
auto *decl = defaultArgument->getParamDecl();
135+
if (decl->hasDefaultExpr()) {
136+
argExpr = decl->getTypeCheckedDefaultExpr();
137+
}
138+
}
139+
parameters.push_back({label, type, extractCompileTimeValue(argExpr)});
140+
}
141+
142+
return parameters;
143+
}
144+
145+
static llvm::Optional<std::string> extractRawLiteral(Expr *expr) {
124146
if (expr) {
125147
switch (expr->getKind()) {
126148
case ExprKind::BooleanLiteral:
@@ -131,7 +153,7 @@ static std::shared_ptr<CompileTimeValue> extractCompileTimeValue(Expr *expr) {
131153
llvm::raw_string_ostream OutputStream(literalOutput);
132154
expr->printConstExprValue(&OutputStream, nullptr);
133155
if (!literalOutput.empty()) {
134-
return std::make_shared<RawLiteralValue>(literalOutput);
156+
return literalOutput;
135157
}
136158
break;
137159
}
@@ -141,7 +163,30 @@ static std::shared_ptr<CompileTimeValue> extractCompileTimeValue(Expr *expr) {
141163
std::string literalOutput;
142164
llvm::raw_string_ostream OutputStream(literalOutput);
143165
OutputStream << stringLiteralExpression->getValue();
144-
return std::make_shared<RawLiteralValue>(literalOutput);
166+
return literalOutput;
167+
}
168+
169+
default:
170+
break;
171+
}
172+
}
173+
return None;
174+
}
175+
176+
static std::shared_ptr<CompileTimeValue> extractCompileTimeValue(Expr *expr) {
177+
if (expr) {
178+
switch (expr->getKind()) {
179+
case ExprKind::BooleanLiteral:
180+
case ExprKind::FloatLiteral:
181+
case ExprKind::IntegerLiteral:
182+
case ExprKind::NilLiteral:
183+
case ExprKind::StringLiteral: {
184+
auto rawLiteral = extractRawLiteral(expr);
185+
if (rawLiteral.hasValue()) {
186+
return std::make_shared<RawLiteralValue>(rawLiteral.value());
187+
}
188+
189+
break;
145190
}
146191

147192
case ExprKind::Array: {
@@ -195,23 +240,38 @@ static std::shared_ptr<CompileTimeValue> extractCompileTimeValue(Expr *expr) {
195240
case ExprKind::Call: {
196241
auto callExpr = cast<CallExpr>(expr);
197242
if (callExpr->getFn()->getKind() == ExprKind::ConstructorRefCall) {
198-
std::vector<FunctionParameter> parameters;
199-
const auto args = callExpr->getArgs();
200-
for (auto arg : *args) {
201-
auto argExpr = arg.getExpr();
202-
const auto label = arg.getLabel().str().str();
203-
const auto type = argExpr->getType();
204-
if (auto defaultArgument = dyn_cast<DefaultArgumentExpr>(argExpr)) {
205-
auto *decl = defaultArgument->getParamDecl();
206-
if (decl->hasDefaultExpr()) {
207-
argExpr = decl->getTypeCheckedDefaultExpr();
208-
}
209-
}
210-
parameters.push_back({label, type, extractCompileTimeValue(argExpr)});
243+
std::vector<FunctionParameter> parameters =
244+
extractFunctionArguments(callExpr->getArgs());
245+
return std::make_shared<InitCallValue>(callExpr->getType(), parameters);
246+
}
247+
248+
if (callExpr->getFn()->getKind() == ExprKind::DotSyntaxCall) {
249+
auto dotSyntaxCallExpr = cast<DotSyntaxCallExpr>(callExpr->getFn());
250+
auto fn = dotSyntaxCallExpr->getFn();
251+
if (fn->getKind() == ExprKind::DeclRef) {
252+
auto declRefExpr = cast<DeclRefExpr>(fn);
253+
auto caseName =
254+
declRefExpr->getDecl()->getName().getBaseIdentifier().str().str();
255+
256+
std::vector<FunctionParameter> parameters =
257+
extractFunctionArguments(callExpr->getArgs());
258+
return std::make_shared<EnumValue>(caseName, parameters);
211259
}
212-
auto name = toFullyQualifiedTypeNameString(callExpr->getType());
213-
return std::make_shared<InitCallValue>(name, parameters);
214260
}
261+
262+
break;
263+
}
264+
265+
case ExprKind::DotSyntaxCall: {
266+
auto dotSyntaxCallExpr = cast<DotSyntaxCallExpr>(expr);
267+
auto fn = dotSyntaxCallExpr->getFn();
268+
if (fn->getKind() == ExprKind::DeclRef) {
269+
auto declRefExpr = cast<DeclRefExpr>(fn);
270+
auto caseName =
271+
declRefExpr->getDecl()->getName().getBaseIdentifier().str().str();
272+
return std::make_shared<EnumValue>(caseName, None);
273+
}
274+
215275
break;
216276
}
217277

@@ -225,11 +285,18 @@ static std::shared_ptr<CompileTimeValue> extractCompileTimeValue(Expr *expr) {
225285
return extractCompileTimeValue(parenExpr->getSubExpr());
226286
}
227287

288+
case ExprKind::PropertyWrapperValuePlaceholder: {
289+
auto placeholderExpr = cast<PropertyWrapperValuePlaceholderExpr>(expr);
290+
return extractCompileTimeValue(
291+
placeholderExpr->getOriginalWrappedValue());
292+
}
293+
228294
default: {
229295
break;
230296
}
231297
}
232298
}
299+
233300
return std::make_shared<RuntimeValue>();
234301
}
235302

@@ -264,7 +331,7 @@ extractCustomAttrValues(VarDecl *propertyDecl) {
264331
static ConstValueTypePropertyInfo
265332
extractTypePropertyInfo(VarDecl *propertyDecl) {
266333
if (const auto binding = propertyDecl->getParentPatternBinding()) {
267-
if (const auto originalInit = binding->getOriginalInit(0)) {
334+
if (const auto originalInit = binding->getInit(0)) {
268335
if (propertyDecl->hasAttachedPropertyWrapper()) {
269336
return {propertyDecl, extractCompileTimeValue(originalInit),
270337
extractCustomAttrValues(propertyDecl)};
@@ -287,6 +354,42 @@ extractTypePropertyInfo(VarDecl *propertyDecl) {
287354
return {propertyDecl, std::make_shared<RuntimeValue>()};
288355
}
289356

357+
llvm::Optional<std::vector<EnumElementDeclValue>>
358+
extractEnumCases(NominalTypeDecl *Decl) {
359+
if (Decl->getKind() == DeclKind::Enum) {
360+
std::vector<EnumElementDeclValue> Elements;
361+
for (EnumCaseDecl *ECD : cast<EnumDecl>(Decl)->getAllCases()) {
362+
for (EnumElementDecl *EED : ECD->getElements()) {
363+
std::string Name = EED->getNameStr().str();
364+
llvm::Optional<std::string> RawValue =
365+
extractRawLiteral(EED->getRawValueExpr());
366+
367+
std::vector<EnumElementParameterValue> Parameters;
368+
if (const ParameterList *Params = EED->getParameterList()) {
369+
for (const ParamDecl *Parameter : Params->getArray()) {
370+
Optional<std::string> Label =
371+
Parameter->getParameterName().empty()
372+
? Optional<std::string>()
373+
: Optional<std::string>(
374+
Parameter->getParameterName().str().str());
375+
376+
Parameters.push_back({Label, Parameter->getType()});
377+
}
378+
}
379+
380+
if (Parameters.empty()) {
381+
Elements.push_back({Name, RawValue, None});
382+
} else {
383+
Elements.push_back({Name, RawValue, Parameters});
384+
}
385+
}
386+
}
387+
return Elements;
388+
}
389+
390+
return None;
391+
}
392+
290393
ConstValueTypeInfo
291394
ConstantValueInfoRequest::evaluate(Evaluator &Evaluator,
292395
NominalTypeDecl *Decl) const {
@@ -317,7 +420,7 @@ ConstantValueInfoRequest::evaluate(Evaluator &Evaluator,
317420
}
318421
}
319422

320-
return ConstValueTypeInfo{Decl, Properties};
423+
return ConstValueTypeInfo{Decl, Properties, extractEnumCases(Decl)};
321424
}
322425

323426
std::vector<ConstValueTypeInfo>
@@ -392,7 +495,8 @@ void writeValue(llvm::json::OStream &JSON,
392495

393496
JSON.attribute("valueKind", "InitCall");
394497
JSON.attributeObject("value", [&]() {
395-
JSON.attribute("type", initCallValue->getName());
498+
JSON.attribute("type",
499+
toFullyQualifiedTypeNameString(initCallValue->getType()));
396500
JSON.attributeArray("arguments", [&] {
397501
for (auto FP : initCallValue->getParameters()) {
398502
JSON.object([&] {
@@ -457,6 +561,27 @@ void writeValue(llvm::json::OStream &JSON,
457561
break;
458562
}
459563

564+
case CompileTimeValue::ValueKind::Enum: {
565+
auto enumValue = cast<EnumValue>(value);
566+
JSON.attribute("valueKind", "Enum");
567+
JSON.attributeObject("value", [&]() {
568+
JSON.attribute("name", enumValue->getIdentifier());
569+
if (enumValue->getParameters().hasValue()) {
570+
auto params = enumValue->getParameters().value();
571+
JSON.attributeArray("arguments", [&] {
572+
for (auto FP : params) {
573+
JSON.object([&] {
574+
JSON.attribute("label", FP.Label);
575+
JSON.attribute("type", toFullyQualifiedTypeNameString(FP.Type));
576+
writeValue(JSON, FP.Value);
577+
});
578+
}
579+
});
580+
}
581+
});
582+
break;
583+
}
584+
460585
case CompileTimeValue::ValueKind::Runtime: {
461586
JSON.attribute("valueKind", "Runtime");
462587
break;
@@ -489,6 +614,38 @@ void writeAttributes(
489614
});
490615
}
491616

617+
void writeEnumCases(
618+
llvm::json::OStream &JSON,
619+
llvm::Optional<std::vector<EnumElementDeclValue>> EnumElements) {
620+
if (!EnumElements.hasValue()) {
621+
return;
622+
}
623+
624+
JSON.attributeArray("cases", [&] {
625+
for (const auto &Case : EnumElements.value()) {
626+
JSON.object([&] {
627+
JSON.attribute("name", Case.Name);
628+
if (Case.RawValue.hasValue()) {
629+
JSON.attribute("rawValue", Case.RawValue.value());
630+
}
631+
if (Case.Parameters.hasValue()) {
632+
JSON.attributeArray("parameters", [&] {
633+
for (const auto &Parameter : Case.Parameters.value()) {
634+
JSON.object([&] {
635+
if (auto Label = Parameter.Label) {
636+
JSON.attribute("label", Label);
637+
}
638+
JSON.attribute("type",
639+
toFullyQualifiedTypeNameString(Parameter.Type));
640+
});
641+
}
642+
});
643+
}
644+
});
645+
}
646+
});
647+
}
648+
492649
bool writeAsJSONToFile(const std::vector<ConstValueTypeInfo> &ConstValueInfos,
493650
llvm::raw_fd_ostream &OS) {
494651
llvm::json::OStream JSON(OS, 2);
@@ -518,6 +675,7 @@ bool writeAsJSONToFile(const std::vector<ConstValueTypeInfo> &ConstValueInfos,
518675
});
519676
}
520677
});
678+
writeEnumCases(JSON, TypeInfo.EnumElements);
521679
});
522680
}
523681
});

0 commit comments

Comments
 (0)