Fixed bug with stack saving

In certain cases when we need to remove the StackSave (because we crossed the save threshold with an allocation), there may already be restores using that stack save which need to be removed as well.

IDEHelper/Backend/BeIRCodeGen.cpp

@@ -1767,8 +1767,11 @@ void BeIRCodeGen::HandleNextCmd()
 			CMD_PARAM(BeValue*, instVal);
 
 			BeInst* inst = (BeInst*)instVal;
-			auto itr = std::find(inst->mParentBlock->mInstructions.begin(), inst->mParentBlock->mInstructions.end(), inst);
+			bool wasRemoved = inst->mParentBlock->mInstructions.Remove(inst);
-			inst->mParentBlock->mInstructions.erase(itr);
+			BF_ASSERT(wasRemoved);
+#ifdef _DEBUG
+			inst->mWasRemoved = true;			
+#endif
 		}
 		break;
 	case BfIRCmd_CreateBr:
@@ -3050,6 +3053,7 @@ BeValue* BeIRCodeGen::GetBeValue(int id)
 	BF_ASSERT(result.mKind == BeIRCodeGenEntryKind_Value);
 #ifdef _DEBUG
 	BF_ASSERT(!result.mBeValue->mLifetimeEnded);
+	BF_ASSERT(!result.mBeValue->mWasRemoved);
 #endif
 	return result.mBeValue;
 }

IDEHelper/Backend/BeModule.cpp

@@ -2940,10 +2940,12 @@ void BeModule::AddBlock(BeFunction* function, BeBlock* block)
 
 void BeModule::RemoveBlock(BeFunction* function, BeBlock* block)
 {
-	auto itr = std::find(function->mBlocks.begin(), function->mBlocks.end(), block);
+	bool didRemove = function->mBlocks.Remove(block);	
-	BF_ASSERT(itr != function->mBlocks.end());
+	BF_ASSERT(didRemove);
-	if (itr != function->mBlocks.end())
+#ifdef _DEBUG
-		function->mBlocks.erase(itr);
+	for (auto inst : block->mInstructions)
+		inst->mWasRemoved = true;
+#endif
 }
 
 BeBlock* BeModule::GetInsertBlock()

IDEHelper/Backend/BeModule.h

@@ -207,9 +207,11 @@ class BeValue : public BeHashble
 public:
 #ifdef _DEBUG
 	bool mLifetimeEnded;
+	bool mWasRemoved;
 	BeValue()
 	{
 		mLifetimeEnded = false;
+		mWasRemoved = false;
 	}
 #endif
 

IDEHelper/Compiler/BfIRBuilder.cpp

@@ -3843,10 +3843,11 @@ BfIRValue BfIRBuilder::CreateStackSave()
 	return retVal;
 }
 
-void BfIRBuilder::CreateStackRestore(BfIRValue stackVal)
+BfIRValue BfIRBuilder::CreateStackRestore(BfIRValue stackVal)
 {
 	BfIRValue retVal = WriteCmd(BfIRCmd_StackRestore, stackVal);
 	NEW_CMD_INSERTED;
+	return retVal;
 }
 
 BfIRValue BfIRBuilder::CreateGlobalVariable(BfIRType varType, bool isConstant, BfIRLinkageType linkageType, BfIRValue initializer, const StringImpl& name, bool isTLS)

IDEHelper/Compiler/BfIRBuilder.h

@@ -1099,7 +1099,7 @@ public:
 	BfIRValue CreateMemSet(BfIRValue addr, BfIRValue val, BfIRValue size, int align);
 	void CreateFence(BfIRFenceType fenceType);
 	BfIRValue CreateStackSave();
-	void CreateStackRestore(BfIRValue stackVal);
+	BfIRValue CreateStackRestore(BfIRValue stackVal);
 	
 	BfIRValue CreateGlobalVariable(BfIRType varType, bool isConstant, BfIRLinkageType linkageType, BfIRValue initializer, const StringImpl& name, bool isTLS = false);	
 	void GlobalVar_SetUnnamedAddr(BfIRValue val, bool unnamedAddr);

IDEHelper/Compiler/BfIRCodeGen.cpp

@@ -1585,6 +1585,7 @@ void BfIRCodeGen::HandleNextCmd()
 			CMD_PARAM(llvm::Value*, stackVal);
 			auto intrin = llvm::Intrinsic::getDeclaration(mLLVMModule, llvm::Intrinsic::stackrestore);
 			auto callInst = mIRBuilder->CreateCall(intrin, llvm::SmallVector<llvm::Value*, 1> {stackVal });			
+			SetResult(curId, callInst);
 		}
 		break;
 	case BfIRCmd_GlobalVariable:

IDEHelper/Compiler/BfModule.cpp

@@ -2122,6 +2122,10 @@ void BfModule::SaveStackState(BfScopeData* scopeData)
 		}
 		if (checkScope->mSavedStack)
 		{
+			for (auto usage : checkScope->mSavedStackUses)
+				mBfIRBuilder->EraseInstFromParent(usage);
+			checkScope->mSavedStackUses.Clear();
+
 			mBfIRBuilder->EraseInstFromParent(checkScope->mSavedStack);
 			checkScope->mSavedStack = BfIRValue();			
 		}		
@@ -12394,7 +12398,7 @@ void BfModule::EmitDeferredScopeCalls(bool useSrcPositions, BfScopeData* scopeDa
 
 			if (checkScope->mSavedStack)
 			{
-				mBfIRBuilder->CreateStackRestore(checkScope->mSavedStack);
+				checkScope->mSavedStackUses.Add(mBfIRBuilder->CreateStackRestore(checkScope->mSavedStack));				
 
 				if (mCurMethodState->mDynStackRevIdx)
 				{

IDEHelper/Compiler/BfModule.h

@@ -315,6 +315,7 @@ public:
 	BfIRValue mBlock;
 	BfIRValue mValueScopeStart;
 	BfIRValue mSavedStack;	
+	Array<BfIRValue> mSavedStackUses;
 	Array<BfDeferredHandler> mDeferredHandlers; // These get cleared when us our a parent gets new entries added into mDeferredCallEntries
 	Array<BfIRBlock> mAtEndBlocks; // Move these to the end after we close scope
 	Array<BfIRValue> mDeferredLifetimeEnds;