llvm/lib/Transforms/Coroutines/CoroInternal.h - llvm-project - Git at Google

 //===- CoroInternal.h - Internal Coroutine interfaces ---------*- C++ -*---===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 // Common definitions/declarations used internally by coroutine lowering passes.
 //===----------------------------------------------------------------------===//

 #ifndef LLVM_LIB_TRANSFORMS_COROUTINES_COROINTERNAL_H
 #define LLVM_LIB_TRANSFORMS_COROUTINES_COROINTERNAL_H

 #include "CoroInstr.h"
 #include "llvm/IR/IRBuilder.h"

 namespace llvm {

 class CallGraph;

 namespace coro {

 bool declaresAnyIntrinsic(const Module &M);
 bool declaresIntrinsics(const Module &M,
                         const std::initializer_list<StringRef>);
 void replaceCoroFree(CoroIdInst *CoroId, bool Elide);

 /// Attempts to rewrite the location operand of debug intrinsics in terms of
 /// the coroutine frame pointer, folding pointer offsets into the DIExpression
 /// of the intrinsic.
 /// If the frame pointer is an Argument, store it into an alloca if
 /// OptimizeFrame is false.
 void salvageDebugInfo(
     SmallDenseMap<Argument *, AllocaInst *, 4> &ArgToAllocaMap,
     DbgVariableIntrinsic &DVI, bool OptimizeFrame, bool IsEntryPoint);
 void salvageDebugInfo(
     SmallDenseMap<Argument *, AllocaInst *, 4> &ArgToAllocaMap, DPValue &DPV,
     bool OptimizeFrame, bool UseEntryValue);

 // Keeps data and helper functions for lowering coroutine intrinsics.
 struct LowererBase {
   Module &TheModule;
   LLVMContext &Context;
   PointerType *const Int8Ptr;
   FunctionType *const ResumeFnType;
   ConstantPointerNull *const NullPtr;

   LowererBase(Module &M);
   Value *makeSubFnCall(Value *Arg, int Index, Instruction *InsertPt);
 };

 enum class ABI {
   /// The "resume-switch" lowering, where there are separate resume and
   /// destroy functions that are shared between all suspend points.  The
   /// coroutine frame implicitly stores the resume and destroy functions,
   /// the current index, and any promise value.
   Switch,

   /// The "returned-continuation" lowering, where each suspend point creates a
   /// single continuation function that is used for both resuming and
   /// destroying.  Does not support promises.
   Retcon,

   /// The "unique returned-continuation" lowering, where each suspend point
   /// creates a single continuation function that is used for both resuming
   /// and destroying.  Does not support promises.  The function is known to
   /// suspend at most once during its execution, and the return value of
   /// the continuation is void.
   RetconOnce,

   /// The "async continuation" lowering, where each suspend point creates a
   /// single continuation function. The continuation function is available as an
   /// intrinsic.
   Async,
 };

 // Holds structural Coroutine Intrinsics for a particular function and other
 // values used during CoroSplit pass.
 struct LLVM_LIBRARY_VISIBILITY Shape {
   CoroBeginInst *CoroBegin;
   SmallVector<AnyCoroEndInst *, 4> CoroEnds;
   SmallVector<CoroSizeInst *, 2> CoroSizes;
   SmallVector<CoroAlignInst *, 2> CoroAligns;
   SmallVector<AnyCoroSuspendInst *, 4> CoroSuspends;
   SmallVector<CallInst*, 2> SwiftErrorOps;

   // Field indexes for special fields in the switch lowering.
   struct SwitchFieldIndex {
     enum {
       Resume,
       Destroy

       // The promise field is always at a fixed offset from the start of
       // frame given its type, but the index isn't a constant for all
       // possible frames.

       // The switch-index field isn't at a fixed offset or index, either;
       // we just work it in where it fits best.
     };
   };

   coro::ABI ABI;

   StructType *FrameTy;
   Align FrameAlign;
   uint64_t FrameSize;
   Value *FramePtr;
   BasicBlock *AllocaSpillBlock;

   /// This would only be true if optimization are enabled.
   bool OptimizeFrame;

   struct SwitchLoweringStorage {
     SwitchInst *ResumeSwitch;
     AllocaInst *PromiseAlloca;
     BasicBlock *ResumeEntryBlock;
     unsigned IndexField;
     unsigned IndexAlign;
     unsigned IndexOffset;
     bool HasFinalSuspend;
     bool HasUnwindCoroEnd;
   };

   struct RetconLoweringStorage {
     Function *ResumePrototype;
     Function *Alloc;
     Function *Dealloc;
     BasicBlock *ReturnBlock;
     bool IsFrameInlineInStorage;
   };

   struct AsyncLoweringStorage {
     Value *Context;
     CallingConv::ID AsyncCC;
     unsigned ContextArgNo;
     uint64_t ContextHeaderSize;
     uint64_t ContextAlignment;
     uint64_t FrameOffset; // Start of the frame.
     uint64_t ContextSize; // Includes frame size.
     GlobalVariable *AsyncFuncPointer;

     Align getContextAlignment() const { return Align(ContextAlignment); }
   };

   union {
     SwitchLoweringStorage SwitchLowering;
     RetconLoweringStorage RetconLowering;
     AsyncLoweringStorage AsyncLowering;
   };

   CoroIdInst *getSwitchCoroId() const {
     assert(ABI == coro::ABI::Switch);
     return cast<CoroIdInst>(CoroBegin->getId());
   }

   AnyCoroIdRetconInst *getRetconCoroId() const {
     assert(ABI == coro::ABI::Retcon ||
            ABI == coro::ABI::RetconOnce);
     return cast<AnyCoroIdRetconInst>(CoroBegin->getId());
   }

   CoroIdAsyncInst *getAsyncCoroId() const {
     assert(ABI == coro::ABI::Async);
     return cast<CoroIdAsyncInst>(CoroBegin->getId());
   }

   unsigned getSwitchIndexField() const {
     assert(ABI == coro::ABI::Switch);
     assert(FrameTy && "frame type not assigned");
     return SwitchLowering.IndexField;
   }
   IntegerType *getIndexType() const {
     assert(ABI == coro::ABI::Switch);
     assert(FrameTy && "frame type not assigned");
     return cast<IntegerType>(FrameTy->getElementType(getSwitchIndexField()));
   }
   ConstantInt *getIndex(uint64_t Value) const {
     return ConstantInt::get(getIndexType(), Value);
   }

   PointerType *getSwitchResumePointerType() const {
     assert(ABI == coro::ABI::Switch);
   assert(FrameTy && "frame type not assigned");
   return cast<PointerType>(FrameTy->getElementType(SwitchFieldIndex::Resume));
   }

   FunctionType *getResumeFunctionType() const {
     switch (ABI) {
     case coro::ABI::Switch:
       return FunctionType::get(Type::getVoidTy(FrameTy->getContext()),
                                PointerType::getUnqual(FrameTy->getContext()),
                                /*IsVarArg=*/false);
     case coro::ABI::Retcon:
     case coro::ABI::RetconOnce:
       return RetconLowering.ResumePrototype->getFunctionType();
     case coro::ABI::Async:
       // Not used. The function type depends on the active suspend.
       return nullptr;
     }

     llvm_unreachable("Unknown coro::ABI enum");
   }

   ArrayRef<Type*> getRetconResultTypes() const {
     assert(ABI == coro::ABI::Retcon ||
            ABI == coro::ABI::RetconOnce);
     auto FTy = CoroBegin->getFunction()->getFunctionType();

     // The safety of all this is checked by checkWFRetconPrototype.
     if (auto STy = dyn_cast<StructType>(FTy->getReturnType())) {
       return STy->elements().slice(1);
     } else {
       return ArrayRef<Type*>();
     }
   }

   ArrayRef<Type*> getRetconResumeTypes() const {
     assert(ABI == coro::ABI::Retcon ||
            ABI == coro::ABI::RetconOnce);

     // The safety of all this is checked by checkWFRetconPrototype.
     auto FTy = RetconLowering.ResumePrototype->getFunctionType();
     return FTy->params().slice(1);
   }

   CallingConv::ID getResumeFunctionCC() const {
     switch (ABI) {
     case coro::ABI::Switch:
       return CallingConv::Fast;

     case coro::ABI::Retcon:
     case coro::ABI::RetconOnce:
       return RetconLowering.ResumePrototype->getCallingConv();
     case coro::ABI::Async:
       return AsyncLowering.AsyncCC;
     }
     llvm_unreachable("Unknown coro::ABI enum");
   }

   AllocaInst *getPromiseAlloca() const {
     if (ABI == coro::ABI::Switch)
       return SwitchLowering.PromiseAlloca;
     return nullptr;
   }

   BasicBlock::iterator getInsertPtAfterFramePtr() const {
     if (auto *I = dyn_cast<Instruction>(FramePtr)) {
       BasicBlock::iterator It = std::next(I->getIterator());
       It.setHeadBit(true); // Copy pre-RemoveDIs behaviour.
       return It;
     }
     return cast<Argument>(FramePtr)->getParent()->getEntryBlock().begin();
   }

   /// Allocate memory according to the rules of the active lowering.
   ///
   /// \param CG - if non-null, will be updated for the new call
   Value *emitAlloc(IRBuilder<> &Builder, Value *Size, CallGraph *CG) const;

   /// Deallocate memory according to the rules of the active lowering.
   ///
   /// \param CG - if non-null, will be updated for the new call
   void emitDealloc(IRBuilder<> &Builder, Value *Ptr, CallGraph *CG) const;

   Shape() = default;
   explicit Shape(Function &F, bool OptimizeFrame = false)
       : OptimizeFrame(OptimizeFrame) {
     buildFrom(F);
   }
   void buildFrom(Function &F);
 };

 bool defaultMaterializable(Instruction &V);
 void buildCoroutineFrame(
     Function &F, Shape &Shape,
     const std::function<bool(Instruction &)> &MaterializableCallback);
 CallInst *createMustTailCall(DebugLoc Loc, Function *MustTailCallFn,
                              ArrayRef<Value *> Arguments, IRBuilder<> &);
 } // End namespace coro.
 } // End namespace llvm

 #endif
	//===- CoroInternal.h - Internal Coroutine interfaces ---------- C++ ----===//
	//
	// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
	// See https://llvm.org/LICENSE.txt for license information.
	// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
	//
	//===----------------------------------------------------------------------===//
	// Common definitions/declarations used internally by coroutine lowering passes.
	//===----------------------------------------------------------------------===//

	#ifndef LLVM_LIB_TRANSFORMS_COROUTINES_COROINTERNAL_H
	#define LLVM_LIB_TRANSFORMS_COROUTINES_COROINTERNAL_H

	#include "CoroInstr.h"
	#include "llvm/IR/IRBuilder.h"

	namespace llvm {

	class CallGraph;

	namespace coro {

	bool declaresAnyIntrinsic(const Module &M);
	bool declaresIntrinsics(const Module &M,
	const std::initializer_list<StringRef>);
	void replaceCoroFree(CoroIdInst *CoroId, bool Elide);

	/// Attempts to rewrite the location operand of debug intrinsics in terms of
	/// the coroutine frame pointer, folding pointer offsets into the DIExpression
	/// of the intrinsic.
	/// If the frame pointer is an Argument, store it into an alloca if
	/// OptimizeFrame is false.
	void salvageDebugInfo(
	SmallDenseMap<Argument , AllocaInst , 4> &ArgToAllocaMap,
	DbgVariableIntrinsic &DVI, bool OptimizeFrame, bool IsEntryPoint);
	void salvageDebugInfo(
	SmallDenseMap<Argument , AllocaInst , 4> &ArgToAllocaMap, DPValue &DPV,
	bool OptimizeFrame, bool UseEntryValue);

	// Keeps data and helper functions for lowering coroutine intrinsics.
	struct LowererBase {
	Module &TheModule;
	LLVMContext &Context;
	PointerType *const Int8Ptr;
	FunctionType *const ResumeFnType;
	ConstantPointerNull *const NullPtr;

	LowererBase(Module &M);
	Value makeSubFnCall(Value Arg, int Index, Instruction *InsertPt);
	};

	enum class ABI {
	/// The "resume-switch" lowering, where there are separate resume and
	/// destroy functions that are shared between all suspend points. The
	/// coroutine frame implicitly stores the resume and destroy functions,
	/// the current index, and any promise value.
	Switch,

	/// The "returned-continuation" lowering, where each suspend point creates a
	/// single continuation function that is used for both resuming and
	/// destroying. Does not support promises.
	Retcon,

	/// The "unique returned-continuation" lowering, where each suspend point
	/// creates a single continuation function that is used for both resuming
	/// and destroying. Does not support promises. The function is known to
	/// suspend at most once during its execution, and the return value of
	/// the continuation is void.
	RetconOnce,

	/// The "async continuation" lowering, where each suspend point creates a
	/// single continuation function. The continuation function is available as an
	/// intrinsic.
	Async,
	};

	// Holds structural Coroutine Intrinsics for a particular function and other
	// values used during CoroSplit pass.
	struct LLVM_LIBRARY_VISIBILITY Shape {
	CoroBeginInst *CoroBegin;
	SmallVector<AnyCoroEndInst *, 4> CoroEnds;
	SmallVector<CoroSizeInst *, 2> CoroSizes;
	SmallVector<CoroAlignInst *, 2> CoroAligns;
	SmallVector<AnyCoroSuspendInst *, 4> CoroSuspends;
	SmallVector<CallInst*, 2> SwiftErrorOps;

	// Field indexes for special fields in the switch lowering.
	struct SwitchFieldIndex {
	enum {
	Resume,
	Destroy

	// The promise field is always at a fixed offset from the start of
	// frame given its type, but the index isn't a constant for all
	// possible frames.

	// The switch-index field isn't at a fixed offset or index, either;
	// we just work it in where it fits best.
	};
	};

	coro::ABI ABI;

	StructType *FrameTy;
	Align FrameAlign;
	uint64_t FrameSize;
	Value *FramePtr;
	BasicBlock *AllocaSpillBlock;

	/// This would only be true if optimization are enabled.
	bool OptimizeFrame;

	struct SwitchLoweringStorage {
	SwitchInst *ResumeSwitch;
	AllocaInst *PromiseAlloca;
	BasicBlock *ResumeEntryBlock;
	unsigned IndexField;
	unsigned IndexAlign;
	unsigned IndexOffset;
	bool HasFinalSuspend;
	bool HasUnwindCoroEnd;
	};

	struct RetconLoweringStorage {
	Function *ResumePrototype;
	Function *Alloc;
	Function *Dealloc;
	BasicBlock *ReturnBlock;
	bool IsFrameInlineInStorage;
	};

	struct AsyncLoweringStorage {
	Value *Context;
	CallingConv::ID AsyncCC;
	unsigned ContextArgNo;
	uint64_t ContextHeaderSize;
	uint64_t ContextAlignment;
	uint64_t FrameOffset; // Start of the frame.
	uint64_t ContextSize; // Includes frame size.
	GlobalVariable *AsyncFuncPointer;

	Align getContextAlignment() const { return Align(ContextAlignment); }
	};

	union {
	SwitchLoweringStorage SwitchLowering;
	RetconLoweringStorage RetconLowering;
	AsyncLoweringStorage AsyncLowering;
	};

	CoroIdInst *getSwitchCoroId() const {
	assert(ABI == coro::ABI::Switch);
	return cast<CoroIdInst>(CoroBegin->getId());
	}

	AnyCoroIdRetconInst *getRetconCoroId() const {
	assert(ABI == coro::ABI::Retcon \|\|
	ABI == coro::ABI::RetconOnce);
	return cast<AnyCoroIdRetconInst>(CoroBegin->getId());
	}

	CoroIdAsyncInst *getAsyncCoroId() const {
	assert(ABI == coro::ABI::Async);
	return cast<CoroIdAsyncInst>(CoroBegin->getId());
	}

	unsigned getSwitchIndexField() const {
	assert(ABI == coro::ABI::Switch);
	assert(FrameTy && "frame type not assigned");
	return SwitchLowering.IndexField;
	}
	IntegerType *getIndexType() const {
	assert(ABI == coro::ABI::Switch);
	assert(FrameTy && "frame type not assigned");
	return cast<IntegerType>(FrameTy->getElementType(getSwitchIndexField()));
	}
	ConstantInt *getIndex(uint64_t Value) const {
	return ConstantInt::get(getIndexType(), Value);
	}

	PointerType *getSwitchResumePointerType() const {
	assert(ABI == coro::ABI::Switch);
	assert(FrameTy && "frame type not assigned");
	return cast<PointerType>(FrameTy->getElementType(SwitchFieldIndex::Resume));
	}

	FunctionType *getResumeFunctionType() const {
	switch (ABI) {
	case coro::ABI::Switch:
	return FunctionType::get(Type::getVoidTy(FrameTy->getContext()),
	PointerType::getUnqual(FrameTy->getContext()),
	/IsVarArg=/false);
	case coro::ABI::Retcon:
	case coro::ABI::RetconOnce:
	return RetconLowering.ResumePrototype->getFunctionType();
	case coro::ABI::Async:
	// Not used. The function type depends on the active suspend.
	return nullptr;
	}

	llvm_unreachable("Unknown coro::ABI enum");
	}

	ArrayRef<Type*> getRetconResultTypes() const {
	assert(ABI == coro::ABI::Retcon \|\|
	ABI == coro::ABI::RetconOnce);
	auto FTy = CoroBegin->getFunction()->getFunctionType();

	// The safety of all this is checked by checkWFRetconPrototype.
	if (auto STy = dyn_cast<StructType>(FTy->getReturnType())) {
	return STy->elements().slice(1);
	} else {
	return ArrayRef<Type*>();
	}
	}

	ArrayRef<Type*> getRetconResumeTypes() const {
	assert(ABI == coro::ABI::Retcon \|\|
	ABI == coro::ABI::RetconOnce);

	// The safety of all this is checked by checkWFRetconPrototype.
	auto FTy = RetconLowering.ResumePrototype->getFunctionType();
	return FTy->params().slice(1);
	}

	CallingConv::ID getResumeFunctionCC() const {
	switch (ABI) {
	case coro::ABI::Switch:
	return CallingConv::Fast;

	case coro::ABI::Retcon:
	case coro::ABI::RetconOnce:
	return RetconLowering.ResumePrototype->getCallingConv();
	case coro::ABI::Async:
	return AsyncLowering.AsyncCC;
	}
	llvm_unreachable("Unknown coro::ABI enum");
	}

	AllocaInst *getPromiseAlloca() const {
	if (ABI == coro::ABI::Switch)
	return SwitchLowering.PromiseAlloca;
	return nullptr;
	}

	BasicBlock::iterator getInsertPtAfterFramePtr() const {
	if (auto *I = dyn_cast<Instruction>(FramePtr)) {
	BasicBlock::iterator It = std::next(I->getIterator());
	It.setHeadBit(true); // Copy pre-RemoveDIs behaviour.
	return It;
	}
	return cast<Argument>(FramePtr)->getParent()->getEntryBlock().begin();
	}

	/// Allocate memory according to the rules of the active lowering.
	///
	/// \param CG - if non-null, will be updated for the new call
	Value emitAlloc(IRBuilder<> &Builder, Value Size, CallGraph *CG) const;

	/// Deallocate memory according to the rules of the active lowering.
	///
	/// \param CG - if non-null, will be updated for the new call
	void emitDealloc(IRBuilder<> &Builder, Value Ptr, CallGraph CG) const;

	Shape() = default;
	explicit Shape(Function &F, bool OptimizeFrame = false)
	: OptimizeFrame(OptimizeFrame) {
	buildFrom(F);
	}
	void buildFrom(Function &F);
	};

	bool defaultMaterializable(Instruction &V);
	void buildCoroutineFrame(
	Function &F, Shape &Shape,
	const std::function<bool(Instruction &)> &MaterializableCallback);
	CallInst createMustTailCall(DebugLoc Loc, Function MustTailCallFn,
	ArrayRef<Value *> Arguments, IRBuilder<> &);
	} // End namespace coro.
	} // End namespace llvm

	#endif