Skip to content

Commit e47fa7a

Browse files
Refactor MLIRGenImpl to eliminate const_cast usage and improve GenContext handling (#211)
1 parent 836b018 commit e47fa7a

1 file changed

Lines changed: 33 additions & 27 deletions

File tree

tslang/lib/TypeScript/MLIRGen.cpp

Lines changed: 33 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -2713,14 +2713,17 @@ class MLIRGenImpl
27132713
mlir::Location location, StringRef name, NodeArray<TypeNode> typeArguments, bool skipThisParam,
27142714
SmallVector<mlir::Value, 4> &operands, const GenContext &genContext)
27152715
{
2716+
// local copy so the 'this'-type override below stays scoped to this instantiation
2717+
GenContext instantiateGenContext(genContext);
2718+
27162719
auto functionGenericTypeInfo = getGenericFunctionInfoByFullName(name);
27172720
if (functionGenericTypeInfo)
27182721
{
2719-
if (functionGenericTypeInfo->functionDeclaration == SyntaxKind::ArrowFunction
2722+
if (functionGenericTypeInfo->functionDeclaration == SyntaxKind::ArrowFunction
27202723
|| functionGenericTypeInfo->functionDeclaration == SyntaxKind::FunctionExpression)
27212724
{
27222725
// we need to avoid wrong redeclaration of arrow functions (when thisType is provided it will add THIS parameter as first)
2723-
const_cast<GenContext &>(genContext).thisType = nullptr;
2726+
instantiateGenContext.thisType = nullptr;
27242727
}
27252728

27262729
MLIRNamespaceGuard ng(currentNamespace);
@@ -2731,7 +2734,7 @@ class MLIRGenImpl
27312734
auto anyNamedGenericType = IsGeneric::False;
27322735

27332736
// step 1, add type arguments first
2734-
GenContext genericTypeGenContext(genContext);
2737+
GenContext genericTypeGenContext(instantiateGenContext);
27352738
genericTypeGenContext.specialization = true;
27362739
genericTypeGenContext.instantiateSpecializedFunction = true;
27372740
genericTypeGenContext.typeParamsWithArgs = functionGenericTypeInfo->typeParamsWithArgs;
@@ -2743,7 +2746,7 @@ class MLIRGenImpl
27432746
{
27442747
// create typeParamsWithArgs from typeArguments
27452748
auto [result, hasAnyNamedGenericType] = zipTypeParametersWithArguments(
2746-
location, typeParams, typeArguments, genericTypeGenContext.typeParamsWithArgs, genContext);
2749+
location, typeParams, typeArguments, genericTypeGenContext.typeParamsWithArgs, instantiateGenContext);
27472750
if (mlir::failed(result))
27482751
{
27492752
return {mlir::failure(), mlir_ts::FunctionType(), ""};
@@ -2783,7 +2786,7 @@ class MLIRGenImpl
27832786

27842787
if (typeParam.getValue().first->getConstraint())
27852788
{
2786-
auto reason = testConstraint(location, genericTypeGenContext.typeParamsWithArgs, typeParamValue.first, widenType, genContext);
2789+
auto reason = testConstraint(location, genericTypeGenContext.typeParamsWithArgs, typeParamValue.first, widenType, instantiateGenContext);
27872790
if (reason == Reason::Failure)
27882791
{
27892792
LLVM_DEBUG(llvm::dbgs() << "\n!! skip. failed. should be resolved later\n";);
@@ -2847,7 +2850,7 @@ class MLIRGenImpl
28472850
return {mlir::success(), funcType, fullName};
28482851
}
28492852

2850-
if (genContext.allowPartialResolve)
2853+
if (instantiateGenContext.allowPartialResolve)
28512854
{
28522855
return {mlir::success(), mlir_ts::FunctionType(), fullName};
28532856
}
@@ -2871,16 +2874,16 @@ class MLIRGenImpl
28712874
auto opIndex = skipThisParam ? 0 : -1;
28722875
// TODO: this is hack, somehow we have difference between operands and call Operands due to CreateExtentionsFunction call
28732876
// review example raytrace.ts function addLight in getNaturalColor (due to captured params)
2874-
long operandsShift = static_cast<long>(operands.size()) - static_cast<long>(genContext.callOperands.size());
2875-
for (auto [callOpIndex, op] : enumerate(genContext.callOperands))
2877+
long operandsShift = static_cast<long>(operands.size()) - static_cast<long>(instantiateGenContext.callOperands.size());
2878+
for (auto [callOpIndex, op] : enumerate(instantiateGenContext.callOperands))
28762879
{
28772880
opIndex++;
28782881
if (isGenericFunctionReference(op))
28792882
{
28802883
LLVM_DEBUG(llvm::dbgs() << "\n!! delayed arrow func instantiation for func type: "
28812884
<< funcOp.getFunctionType() << "\n";);
28822885
auto result = instantiateSpecializedFunction(
2883-
location, op, funcOp.getFunctionType().getInput(opIndex), genContext);
2886+
location, op, funcOp.getFunctionType().getInput(opIndex), instantiateGenContext);
28842887
if (mlir::failed(result))
28852888
{
28862889
return {mlir::failure(), mlir_ts::FunctionType(), ""};
@@ -8000,27 +8003,28 @@ class MLIRGenImpl
80008003
label = "";
80018004
}
80028005

8003-
const_cast<GenContext &>(genContext).isLoop = true;
8004-
const_cast<GenContext &>(genContext).loopLabel = label;
8006+
GenContext loopGenContext(genContext);
8007+
loopGenContext.isLoop = true;
8008+
loopGenContext.loopLabel = label;
80058009

80068010
/*auto *cond =*/builder.createBlock(&doWhileOp.getCond(), {}, types);
80078011
/*auto *body =*/builder.createBlock(&doWhileOp.getBody(), {}, types);
80088012

80098013
// body in condition
80108014
builder.setInsertionPointToStart(&doWhileOp.getBody().front());
8011-
auto result2 = mlirGen(doStatementAST->statement, genContext);
8015+
auto result2 = mlirGen(doStatementAST->statement, loopGenContext);
80128016
EXIT_IF_FAILED(result2)
80138017
// just simple return, as body in cond
80148018
builder.create<mlir_ts::ResultOp>(location);
80158019

80168020
builder.setInsertionPointToStart(&doWhileOp.getCond().front());
8017-
auto result = mlirGen(doStatementAST->expression, genContext);
8021+
auto result = mlirGen(doStatementAST->expression, loopGenContext);
80188022
EXIT_IF_FAILED(result)
80198023
auto conditionValue = V(result);
80208024

80218025
if (conditionValue.getType() != getBooleanType())
80228026
{
8023-
CAST(conditionValue, location, getBooleanType(), conditionValue, genContext);
8027+
CAST(conditionValue, location, getBooleanType(), conditionValue, loopGenContext);
80248028
}
80258029

80268030
builder.create<mlir_ts::ConditionOp>(location, conditionValue, mlir::ValueRange{});
@@ -8045,21 +8049,22 @@ class MLIRGenImpl
80458049
label = "";
80468050
}
80478051

8048-
const_cast<GenContext &>(genContext).isLoop = true;
8049-
const_cast<GenContext &>(genContext).loopLabel = label;
8052+
GenContext loopGenContext(genContext);
8053+
loopGenContext.isLoop = true;
8054+
loopGenContext.loopLabel = label;
80508055

80518056
/*auto *cond =*/builder.createBlock(&whileOp.getCond(), {}, types);
80528057
/*auto *body =*/builder.createBlock(&whileOp.getBody(), {}, types);
80538058

80548059
// condition
80558060
builder.setInsertionPointToStart(&whileOp.getCond().front());
8056-
auto result = mlirGen(whileStatementAST->expression, genContext);
8061+
auto result = mlirGen(whileStatementAST->expression, loopGenContext);
80578062
EXIT_IF_FAILED_OR_NO_VALUE(result)
80588063
auto conditionValue = V(result);
80598064

80608065
if (conditionValue.getType() != getBooleanType())
80618066
{
8062-
CAST(conditionValue, location, getBooleanType(), conditionValue, genContext);
8067+
CAST(conditionValue, location, getBooleanType(), conditionValue, loopGenContext);
80638068
}
80648069

80658070
builder.create<mlir_ts::ConditionOp>(location, conditionValue, mlir::ValueRange{});
@@ -8070,9 +8075,9 @@ class MLIRGenImpl
80708075
// check if we do safe-cast here
80718076
SymbolTableScopeT varScopeBody(symbolTable);
80728077
SafeTypesMapScopeT safeTypesMapScope(safeTypesMap);
8073-
checkSafeCast(whileStatementAST->expression, conditionValue, nullptr, genContext);
8078+
checkSafeCast(whileStatementAST->expression, conditionValue, nullptr, loopGenContext);
80748079

8075-
auto result2 = mlirGen(whileStatementAST->statement, genContext);
8080+
auto result2 = mlirGen(whileStatementAST->statement, loopGenContext);
80768081
EXIT_IF_FAILED(result2)
80778082
builder.create<mlir_ts::ResultOp>(location);
80788083

@@ -8131,15 +8136,16 @@ class MLIRGenImpl
81318136
label = "";
81328137
}
81338138

8134-
const_cast<GenContext &>(genContext).isLoop = true;
8135-
const_cast<GenContext &>(genContext).loopLabel = label;
8139+
GenContext loopGenContext(genContext);
8140+
loopGenContext.isLoop = true;
8141+
loopGenContext.loopLabel = label;
81368142

81378143
/*auto *cond =*/builder.createBlock(&forOp.getCond(), {}, types);
81388144
/*auto *body =*/builder.createBlock(&forOp.getBody(), {}, types);
81398145
/*auto *incr =*/builder.createBlock(&forOp.getIncr(), {}, types);
81408146

81418147
builder.setInsertionPointToStart(&forOp.getCond().front());
8142-
auto result = mlirGen(forStatementAST->condition, genContext);
8148+
auto result = mlirGen(forStatementAST->condition, loopGenContext);
81438149
EXIT_IF_FAILED(result)
81448150
auto conditionValue = V(result);
81458151
if (conditionValue)
@@ -8158,7 +8164,7 @@ class MLIRGenImpl
81588164
if (forStatementAST->statement == SyntaxKind::Block)
81598165
{
81608166
auto firstStatement = forStatementAST->statement.as<ts::Block>()->statements.front();
8161-
auto result = mlirGen(firstStatement, genContext);
8167+
auto result = mlirGen(firstStatement, loopGenContext);
81628168
EXIT_IF_FAILED(result)
81638169
}
81648170

@@ -8168,7 +8174,7 @@ class MLIRGenImpl
81688174
auto asyncExecOp = builder.create<mlir::async::ExecuteOp>(
81698175
stripMetadata(location), mlir::TypeRange{}, mlir::ValueRange{}, mlir::ValueRange{},
81708176
[&](mlir::OpBuilder &builder, mlir::Location location, mlir::ValueRange values) {
8171-
GenContext execOpBodyGenContext(genContext);
8177+
GenContext execOpBodyGenContext(loopGenContext);
81728178
DITableScopeT debugAsyncCodeScope(debugScope);
81738179
MLIRDebugInfoHelper mdi(builder, debugScope);
81748180

@@ -8207,15 +8213,15 @@ class MLIRGenImpl
82078213
else
82088214
{
82098215
// default
8210-
auto result = mlirGen(forStatementAST->statement, genContext);
8216+
auto result = mlirGen(forStatementAST->statement, loopGenContext);
82118217
EXIT_IF_FAILED(result)
82128218
}
82138219

82148220
builder.create<mlir_ts::ResultOp>(location);
82158221

82168222
// increment
82178223
builder.setInsertionPointToStart(&forOp.getIncr().front());
8218-
mlirGen(forStatementAST->incrementor, genContext);
8224+
mlirGen(forStatementAST->incrementor, loopGenContext);
82198225
builder.create<mlir_ts::ResultOp>(location);
82208226

82218227
builder.setInsertionPointAfter(forOp);

0 commit comments

Comments
 (0)