src/transform/canonicalize_entry_point_io.cc - tint - Git at Google

 // Copyright 2021 The Tint Authors.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
 //     http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.

 #include "src/transform/canonicalize_entry_point_io.h"

 #include <utility>

 #include "src/program_builder.h"
 #include "src/semantic/function.h"
 #include "src/semantic/statement.h"
 #include "src/semantic/variable.h"

 namespace tint {
 namespace transform {

 CanonicalizeEntryPointIO::CanonicalizeEntryPointIO() = default;
 CanonicalizeEntryPointIO::~CanonicalizeEntryPointIO() = default;

 Transform::Output CanonicalizeEntryPointIO::Run(const Program* in,
                                                 const DataMap&) {
   ProgramBuilder out;
   CloneContext ctx(&out, in);

   // Start by cloning all the symbols. This ensures that the authored symbols
   // won't get renamed if they collide with new symbols below.
   ctx.CloneSymbols();

   // Strip entry point IO decorations from struct declarations.
   // TODO(jrprice): This code is duplicated with the SPIR-V transform.
   for (auto* ty : ctx.src->AST().ConstructedTypes()) {
     if (auto* struct_ty = ty->As<type::Struct>()) {
       // Build new list of struct members without entry point IO decorations.
       ast::StructMemberList new_struct_members;
       for (auto* member : struct_ty->impl()->members()) {
         ast::DecorationList new_decorations = RemoveDecorations(
             &ctx, member->decorations(), [](const ast::Decoration* deco) {
               return deco
                   ->IsAnyOf<ast::BuiltinDecoration, ast::LocationDecoration>();
             });
         new_struct_members.push_back(
             ctx.dst->Member(ctx.Clone(member->symbol()),
                             ctx.Clone(member->type()), new_decorations));
       }

       // Redeclare the struct.
       auto* new_struct = ctx.dst->create<type::Struct>(
           ctx.Clone(struct_ty->symbol()),
           ctx.dst->create<ast::Struct>(
               new_struct_members, ctx.Clone(struct_ty->impl()->decorations())));
       ctx.Replace(struct_ty, new_struct);
     }
   }

   for (auto* func : ctx.src->AST().Functions()) {
     if (!func->IsEntryPoint()) {
       continue;
     }

     ast::VariableList new_parameters;

     if (!func->params().empty()) {
       // Collect all parameters and build a list of new struct members.
       auto new_struct_param_symbol = ctx.dst->Symbols().New();
       ast::StructMemberList new_struct_members;
       for (auto* param : func->params()) {
         auto param_name = ctx.Clone(param->symbol());
         auto* param_ty = ctx.src->Sem().Get(param)->Type();
         auto* param_declared_ty = ctx.src->Sem().Get(param)->DeclaredType();

         ast::Expression* func_const_initializer = nullptr;

         if (auto* struct_ty = param_ty->As<type::Struct>()) {
           // Pull out all struct members and build initializer list.
           ast::ExpressionList init_values;
           for (auto* member : struct_ty->impl()->members()) {
             if (member->type()->UnwrapAll()->Is<type::Struct>()) {
               TINT_ICE(ctx.dst->Diagnostics()) << "nested pipeline IO struct";
             }

             ast::DecorationList new_decorations = RemoveDecorations(
                 &ctx, member->decorations(), [](const ast::Decoration* deco) {
                   return !deco->IsAnyOf<ast::BuiltinDecoration,
                                         ast::LocationDecoration>();
                 });
             auto member_name = ctx.Clone(member->symbol());
             new_struct_members.push_back(ctx.dst->Member(
                 member_name, ctx.Clone(member->type()), new_decorations));
             init_values.push_back(
                 ctx.dst->MemberAccessor(new_struct_param_symbol, member_name));
           }

           func_const_initializer =
               ctx.dst->Construct(ctx.Clone(param_declared_ty), init_values);
         } else {
           ast::DecorationList new_decorations = RemoveDecorations(
               &ctx, param->decorations(), [](const ast::Decoration* deco) {
                 return !deco->IsAnyOf<ast::BuiltinDecoration,
                                       ast::LocationDecoration>();
               });
           new_struct_members.push_back(ctx.dst->Member(
               param_name, ctx.Clone(param_declared_ty), new_decorations));
           func_const_initializer =
               ctx.dst->MemberAccessor(new_struct_param_symbol, param_name);
         }

         if (func->body()->empty()) {
           // Don't generate a function-scope const if the function is empty.
           continue;
         }

         // Create a function-scope const to replace the parameter.
         // Initialize it with the value extracted from the new struct parameter.
         auto* func_const = ctx.dst->Const(
             param_name, ctx.Clone(param_declared_ty), func_const_initializer);
         ctx.InsertBefore(func->body()->statements(), *func->body()->begin(),
                          ctx.dst->WrapInStatement(func_const));

         // Replace all uses of the function parameter with the function const.
         for (auto* user : ctx.src->Sem().Get(param)->Users()) {
           ctx.Replace<ast::Expression>(user->Declaration(),
                                        ctx.dst->Expr(param_name));
         }
       }

       // Create the new struct type.
       auto* in_struct = ctx.dst->create<type::Struct>(
           ctx.dst->Symbols().New(),
           ctx.dst->create<ast::Struct>(new_struct_members,
                                        ast::DecorationList{}));
       ctx.InsertBefore(ctx.src->AST().GlobalDeclarations(), func, in_struct);

       // Create a new function parameter using this struct type.
       auto* struct_param = ctx.dst->Param(new_struct_param_symbol, in_struct);
       new_parameters.push_back(struct_param);
     }

     // Handle return type.
     auto* ret_type = func->return_type()->UnwrapAliasIfNeeded();
     type::Type* new_ret_type;
     if (ret_type->Is<type::Void>()) {
       new_ret_type = ctx.dst->ty.void_();
     } else {
       ast::StructMemberList new_struct_members;

       if (auto* struct_ty = ret_type->As<type::Struct>()) {
         // Rebuild struct with only the entry point IO attributes.
         for (auto* member : struct_ty->impl()->members()) {
           if (member->type()->UnwrapAll()->Is<type::Struct>()) {
             TINT_ICE(ctx.dst->Diagnostics()) << "nested pipeline IO struct";
           }

           ast::DecorationList new_decorations = RemoveDecorations(
               &ctx, member->decorations(), [](const ast::Decoration* deco) {
                 return !deco->IsAnyOf<ast::BuiltinDecoration,
                                       ast::LocationDecoration>();
               });
           new_struct_members.push_back(
               ctx.dst->Member(ctx.Clone(member->symbol()),
                               ctx.Clone(member->type()), new_decorations));
         }
       } else {
         new_struct_members.push_back(
             ctx.dst->Member("value", ctx.Clone(ret_type),
                             ctx.Clone(func->return_type_decorations())));
       }

       // Create the new struct type.
       auto* out_struct = ctx.dst->create<type::Struct>(
           ctx.dst->Symbols().New(),
           ctx.dst->create<ast::Struct>(new_struct_members,
                                        ast::DecorationList{}));
       ctx.InsertBefore(ctx.src->AST().GlobalDeclarations(), func, out_struct);
       new_ret_type = out_struct;

       // Replace all return statements.
       auto* sem_func = ctx.src->Sem().Get(func);
       for (auto* ret : sem_func->ReturnStatements()) {
         auto* ret_sem = ctx.src->Sem().Get(ret);
         // Reconstruct the return value using the newly created struct.
         auto* new_ret_value = ctx.Clone(ret->value());
         ast::ExpressionList ret_values;
         if (auto* struct_ty = ret_type->As<type::Struct>()) {
           if (!ret->value()->Is<ast::IdentifierExpression>()) {
             // Create a const to hold the return value expression to avoid
             // re-evaluating it multiple times.
             auto temp = ctx.dst->Symbols().New();
             auto* temp_var = ctx.dst->Decl(
                 ctx.dst->Const(temp, ctx.Clone(ret_type), new_ret_value));
             ctx.InsertBefore(ret_sem->Block()->statements(), ret, temp_var);
             new_ret_value = ctx.dst->Expr(temp);
           }

           for (auto* member : struct_ty->impl()->members()) {
             ret_values.push_back(ctx.dst->MemberAccessor(
                 new_ret_value, ctx.Clone(member->symbol())));
           }
         } else {
           ret_values.push_back(new_ret_value);
         }

         auto* new_ret = ctx.dst->create<ast::ReturnStatement>(
             ctx.dst->Construct(new_ret_type, ret_values));
         ctx.Replace(ret, new_ret);
       }
     }

     // Rewrite the function header with the new parameters.
     auto* new_func = ctx.dst->create<ast::Function>(
         func->source(), ctx.Clone(func->symbol()), new_parameters, new_ret_type,
         ctx.Clone(func->body()), ctx.Clone(func->decorations()),
         ast::DecorationList{});
     ctx.Replace(func, new_func);
   }

   ctx.Clone();
   return Output(Program(std::move(out)));
 }

 }  // namespace transform
 }  // namespace tint
	// Copyright 2021 The Tint Authors.
	//
	// Licensed under the Apache License, Version 2.0 (the "License");
	// you may not use this file except in compliance with the License.
	// You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.

	#include "src/transform/canonicalize_entry_point_io.h"

	#include <utility>

	#include "src/program_builder.h"
	#include "src/semantic/function.h"
	#include "src/semantic/statement.h"
	#include "src/semantic/variable.h"

	namespace tint {
	namespace transform {

	CanonicalizeEntryPointIO::CanonicalizeEntryPointIO() = default;
	CanonicalizeEntryPointIO::~CanonicalizeEntryPointIO() = default;

	Transform::Output CanonicalizeEntryPointIO::Run(const Program* in,
	const DataMap&) {
	ProgramBuilder out;
	CloneContext ctx(&out, in);

	// Start by cloning all the symbols. This ensures that the authored symbols
	// won't get renamed if they collide with new symbols below.
	ctx.CloneSymbols();

	// Strip entry point IO decorations from struct declarations.
	// TODO(jrprice): This code is duplicated with the SPIR-V transform.
	for (auto* ty : ctx.src->AST().ConstructedTypes()) {
	if (auto* struct_ty = ty->As<type::Struct>()) {
	// Build new list of struct members without entry point IO decorations.
	ast::StructMemberList new_struct_members;
	for (auto* member : struct_ty->impl()->members()) {
	ast::DecorationList new_decorations = RemoveDecorations(
	&ctx, member->decorations(), [](const ast::Decoration* deco) {
	return deco
	->IsAnyOf<ast::BuiltinDecoration, ast::LocationDecoration>();
	});
	new_struct_members.push_back(
	ctx.dst->Member(ctx.Clone(member->symbol()),
	ctx.Clone(member->type()), new_decorations));
	}

	// Redeclare the struct.
	auto* new_struct = ctx.dst->create<type::Struct>(
	ctx.Clone(struct_ty->symbol()),
	ctx.dst->create<ast::Struct>(
	new_struct_members, ctx.Clone(struct_ty->impl()->decorations())));
	ctx.Replace(struct_ty, new_struct);
	}
	}

	for (auto* func : ctx.src->AST().Functions()) {
	if (!func->IsEntryPoint()) {
	continue;
	}

	ast::VariableList new_parameters;

	if (!func->params().empty()) {
	// Collect all parameters and build a list of new struct members.
	auto new_struct_param_symbol = ctx.dst->Symbols().New();
	ast::StructMemberList new_struct_members;
	for (auto* param : func->params()) {
	auto param_name = ctx.Clone(param->symbol());
	auto* param_ty = ctx.src->Sem().Get(param)->Type();
	auto* param_declared_ty = ctx.src->Sem().Get(param)->DeclaredType();

	ast::Expression* func_const_initializer = nullptr;

	if (auto* struct_ty = param_ty->As<type::Struct>()) {
	// Pull out all struct members and build initializer list.
	ast::ExpressionList init_values;
	for (auto* member : struct_ty->impl()->members()) {
	if (member->type()->UnwrapAll()->Is<type::Struct>()) {
	TINT_ICE(ctx.dst->Diagnostics()) << "nested pipeline IO struct";
	}

	ast::DecorationList new_decorations = RemoveDecorations(
	&ctx, member->decorations(), [](const ast::Decoration* deco) {
	return !deco->IsAnyOf<ast::BuiltinDecoration,
	ast::LocationDecoration>();
	});
	auto member_name = ctx.Clone(member->symbol());
	new_struct_members.push_back(ctx.dst->Member(
	member_name, ctx.Clone(member->type()), new_decorations));
	init_values.push_back(
	ctx.dst->MemberAccessor(new_struct_param_symbol, member_name));
	}

	func_const_initializer =
	ctx.dst->Construct(ctx.Clone(param_declared_ty), init_values);
	} else {
	ast::DecorationList new_decorations = RemoveDecorations(
	&ctx, param->decorations(), [](const ast::Decoration* deco) {
	return !deco->IsAnyOf<ast::BuiltinDecoration,
	ast::LocationDecoration>();
	});
	new_struct_members.push_back(ctx.dst->Member(
	param_name, ctx.Clone(param_declared_ty), new_decorations));
	func_const_initializer =
	ctx.dst->MemberAccessor(new_struct_param_symbol, param_name);
	}

	if (func->body()->empty()) {
	// Don't generate a function-scope const if the function is empty.
	continue;
	}

	// Create a function-scope const to replace the parameter.
	// Initialize it with the value extracted from the new struct parameter.
	auto* func_const = ctx.dst->Const(
	param_name, ctx.Clone(param_declared_ty), func_const_initializer);
	ctx.InsertBefore(func->body()->statements(), *func->body()->begin(),
	ctx.dst->WrapInStatement(func_const));

	// Replace all uses of the function parameter with the function const.
	for (auto* user : ctx.src->Sem().Get(param)->Users()) {
	ctx.Replace<ast::Expression>(user->Declaration(),
	ctx.dst->Expr(param_name));
	}
	}

	// Create the new struct type.
	auto* in_struct = ctx.dst->create<type::Struct>(
	ctx.dst->Symbols().New(),
	ctx.dst->create<ast::Struct>(new_struct_members,
	ast::DecorationList{}));
	ctx.InsertBefore(ctx.src->AST().GlobalDeclarations(), func, in_struct);

	// Create a new function parameter using this struct type.
	auto* struct_param = ctx.dst->Param(new_struct_param_symbol, in_struct);
	new_parameters.push_back(struct_param);
	}

	// Handle return type.
	auto* ret_type = func->return_type()->UnwrapAliasIfNeeded();
	type::Type* new_ret_type;
	if (ret_type->Is<type::Void>()) {
	new_ret_type = ctx.dst->ty.void_();
	} else {
	ast::StructMemberList new_struct_members;

	if (auto* struct_ty = ret_type->As<type::Struct>()) {
	// Rebuild struct with only the entry point IO attributes.
	for (auto* member : struct_ty->impl()->members()) {
	if (member->type()->UnwrapAll()->Is<type::Struct>()) {
	TINT_ICE(ctx.dst->Diagnostics()) << "nested pipeline IO struct";
	}

	ast::DecorationList new_decorations = RemoveDecorations(
	&ctx, member->decorations(), [](const ast::Decoration* deco) {
	return !deco->IsAnyOf<ast::BuiltinDecoration,
	ast::LocationDecoration>();
	});
	new_struct_members.push_back(
	ctx.dst->Member(ctx.Clone(member->symbol()),
	ctx.Clone(member->type()), new_decorations));
	}
	} else {
	new_struct_members.push_back(
	ctx.dst->Member("value", ctx.Clone(ret_type),
	ctx.Clone(func->return_type_decorations())));
	}

	// Create the new struct type.
	auto* out_struct = ctx.dst->create<type::Struct>(
	ctx.dst->Symbols().New(),
	ctx.dst->create<ast::Struct>(new_struct_members,
	ast::DecorationList{}));
	ctx.InsertBefore(ctx.src->AST().GlobalDeclarations(), func, out_struct);
	new_ret_type = out_struct;

	// Replace all return statements.
	auto* sem_func = ctx.src->Sem().Get(func);
	for (auto* ret : sem_func->ReturnStatements()) {
	auto* ret_sem = ctx.src->Sem().Get(ret);
	// Reconstruct the return value using the newly created struct.
	auto* new_ret_value = ctx.Clone(ret->value());
	ast::ExpressionList ret_values;
	if (auto* struct_ty = ret_type->As<type::Struct>()) {
	if (!ret->value()->Is<ast::IdentifierExpression>()) {
	// Create a const to hold the return value expression to avoid
	// re-evaluating it multiple times.
	auto temp = ctx.dst->Symbols().New();
	auto* temp_var = ctx.dst->Decl(
	ctx.dst->Const(temp, ctx.Clone(ret_type), new_ret_value));
	ctx.InsertBefore(ret_sem->Block()->statements(), ret, temp_var);
	new_ret_value = ctx.dst->Expr(temp);
	}

	for (auto* member : struct_ty->impl()->members()) {
	ret_values.push_back(ctx.dst->MemberAccessor(
	new_ret_value, ctx.Clone(member->symbol())));
	}
	} else {
	ret_values.push_back(new_ret_value);
	}

	auto* new_ret = ctx.dst->create<ast::ReturnStatement>(
	ctx.dst->Construct(new_ret_type, ret_values));
	ctx.Replace(ret, new_ret);
	}
	}

	// Rewrite the function header with the new parameters.
	auto* new_func = ctx.dst->create<ast::Function>(
	func->source(), ctx.Clone(func->symbol()), new_parameters, new_ret_type,
	ctx.Clone(func->body()), ctx.Clone(func->decorations()),
	ast::DecorationList{});
	ctx.Replace(func, new_func);
	}

	ctx.Clone();
	return Output(Program(std::move(out)));
	}

	} // namespace transform
	} // namespace tint