Undo Manager

#ifndef UNDOMGR_H__831D7338_E2BC_4405_BBF2_9FBB761AE1E9
#define UNDOMGR_H__831D7338_E2BC_4405_BBF2_9FBB761AE1E9

//////////////////////////////////////////////////////////////////////////
// Copyright (C) 2001 by Jens Nilsson, jnilsson@icebreaker.com.
//
// This code is free for personal and commercial use, providing this 
// notice remains intact in the source files and all eventual changes are
// clearly marked with comments.
//
// No warranty of any kind, expressed or implied, is included with this
// software; use at your own risk. Responsibility for damages (if any) to
// anyone resulting from the use of this software rests entirely with the
// user.
//////////////////////////////////////////////////////////////////////////

#ifndef _UNDOMGR_H_
#define _UNDOMGR_H_
#endif

#if _MSC_VER >= 1000
#pragma once
#endif // _MSC_VER >= 1000

#ifndef __ATLBASE_H__
	#error undomgr.h requires atlbase.h to be included first
#endif

#include <list>
#include <algorithm>

namespace ATLOLE
{

/*
* Using a list for the undo and redo stacks
*/
typedef std::list<CAdapt<CComPtr<IOleUndoUnit> > > OleUndoUnitList;

/*
* Function object to call Do on child units, used with the 
* for_each algorithm
*/
struct CallDoFunction
{
	IOleUndoManager* m_pManager;
	HRESULT			 m_hr;

	CallDoFunction(IOleUndoManager* pManager = NULL) : m_pManager(pManager), m_hr(S_OK) {}

	void operator()(CAdapt<CComPtr<IOleUndoUnit> > p)
	{
		HRESULT hr = p.m_T->Do(m_pManager);
		if (FAILED(hr))
			m_hr = hr;
	}

	HRESULT GetLastError() 
	{ 
		return m_hr; 
	}
};

/*
* Function object to find a specific unit. Used with the find_if algorithm.
*/
struct FindUnitFunction
{
	IOleUndoUnit* m_pUnit;

	FindUnitFunction(IOleUndoUnit* pUnit) : m_pUnit(pUnit) {}

	bool operator()(CAdapt<CComPtr<IOleUndoUnit> > p)
	{
		if (p.m_T.IsEqualObject(m_pUnit))
			return true;

		// See if this is a parent unit and see if the unit is a child
		CComQIPtr<IOleParentUndoUnit> q = p.m_T;
		if (q != NULL)
			return q->FindUnit(m_pUnit) == S_OK ? true : false;

		return false;
	}
};

/*
* In case we have no CComCoClass, this template imlements GetObjectCLSID
*/
template <const CLSID* pclsid = &CLSID_NULL>
class CComClassID
{
public:
	static const CLSID& WINAPI GetObjectCLSID() {return *pclsid;}
};


/*
* Basic IOleUndoUnit implementation
*/
template<class T, LONG lTypeID=0>
class ATL_NO_VTABLE IOleUndoUnitImpl : public IOleUndoUnit
{
public:

	//To be implemented by derived class
	HRESULT IOleUndoUnitImpl_Do(IOleUndoManager* /*pUndoManager*/)
	{
		// Do work here
		//
		// If Undo operation fail, return E_ABORT and the undo manager
		// will try to rollback the current undo command.
		return S_OK;
	}

	HRESULT IOleUndoUnitImpl_CreateUndoUnit(IOleUndoUnit** /*ppUU*/)
	{
		// Create redo - or an undo unit here
		//
		// If creation fails, return E_ABORT and the undo manager
		// will try to rollback the current undo command.
		ATLASSERT(FALSE);
		return E_NOTIMPL;
	}

	//Interface part
	STDMETHOD(Do)(IOleUndoManager* pUndoManager)
	{
		T* pT = static_cast<T*>(this);
		HRESULT hr = pT->IOleUndoUnitImpl_Do(pUndoManager);
		if (SUCCEEDED(hr))
		{
			if (pUndoManager)
			{
				CComPtr<IOleUndoUnit> spUU;
				HRESULT hRet = pT->IOleUndoUnitImpl_CreateUndoUnit(&spUU);
				if (SUCCEEDED(hRet))
				{
					hr = pUndoManager->Add(spUU);
				}
				else if (hRet != E_NOTIMPL)
				{
					hr = hRet;
				}
			}
		}
		return hr;
	}
	STDMETHOD(GetDescription)(BSTR* pBstr)
	{
		//Needs to be implemented by user
		ATLASSERT(FALSE);
		*pBstr = NULL;
		return E_NOTIMPL;
	}
	STDMETHOD(GetUnitType)(CLSID* pClsid, LONG* plID)
	{
		*pClsid = T::GetObjectCLSID();
		*plID = lTypeID;
		return S_OK;
	}
	STDMETHOD(OnNextAdd)()
	{
		// This method might be called for the last undo unit to notify
		// that a new undo unit has been added.
		// No further data can be fed into this undo unit from now.
		return S_OK;
	}
};


/*
* Basic IOleParentUndoUnit implementation
*/
template<class T, LONG lTypeID=0>
class ATL_NO_VTABLE IOleParentUndoUnitImpl : public IOleParentUndoUnit
{
public:

	CComPtr<IOleParentUndoUnit> m_pOpenParentUndoUnit;
	OleUndoUnitList				m_childUnits;


	//To be implemented by derived classes
	HRESULT IOleParentUndoUnitImpl_Do(IOleUndoManager* /*pUndoManager*/)
	{
		// It is optional to implement this method, typically a parent unit
		// does not perform actions by it self.
		return S_OK;
	}

	HRESULT IOleParentUndoUnitImpl_CreateParentUndoUnit(IOleParentUndoUnit** /*ppPUU*/)
	{
		// If creation fails, return E_ABORT and the undo manager
		// will try to rollback the current undo command.
		ATLASSERT(FALSE);
		return E_NOTIMPL;
	}

	// Overload this method to provide context sensitive state
	DWORD GetState()
	{
		// NOTE: State should not be blocked while a parent unit is open
		return UAS_NORMAL; 
	}

	// Interface part
	STDMETHOD(Do)(IOleUndoManager* pUndoManager)
	{
		T* pT = static_cast<T*>(this);

		bool bOpenedParent = false;
		HRESULT hr = S_OK;
		CComPtr<IOleParentUndoUnit> spPUU;

		if (pUndoManager)
		{
			hr = pT->IOleParentUndoUnitImpl_CreateParentUndoUnit(&spPUU);
			if (SUCCEEDED(hr))
			{
				hr = pUndoManager->Open(spPUU);
				if (SUCCEEDED(hr))
					bOpenedParent = true;
			}
		}

		// Give the parent unit a chance to do its work before calling its cildren
		if (SUCCEEDED(hr))
			hr = IOleParentUndoUnitImpl_Do(pUndoManager);

		if (SUCCEEDED(hr))
			hr = std::for_each( m_childUnits.begin(), 
								m_childUnits.end(), 
								CallDoFunction(pUndoManager)).GetLastError();

		if (pUndoManager && bOpenedParent)
		{
			BOOL bCommit = SUCCEEDED(hr); // If errors occured discard the unit
			HRESULT hRet = pUndoManager->Close(spPUU, bCommit);
			if (bCommit) // Pass along this error if we tried to commit
				hr = hRet;
		}

		return hr;
	}
	STDMETHOD(GetDescription)(BSTR* pBstr)
	{
		//Needs to be implemented by user
		ATLASSERT(FALSE);
		*pBstr = NULL;
		return E_NOTIMPL;
	}
	STDMETHOD(GetUnitType)(CLSID* pClsid, LONG* plID)
	{
		*pClsid = T::GetObjectCLSID();
		*plID = lTypeID;
		return S_OK;
	}
	STDMETHOD(OnNextAdd)()
	{
		// This method might be called to the last undo unit to notify
		// it that a new undo unit has been added.
		if (!m_childUnits.empty())
		{
			(*m_childUnits.begin()).m_T->OnNextAdd();
		}
		return S_OK;
	}
	STDMETHOD(Open)(IOleParentUndoUnit* pPUU)
	{
		T* pT = static_cast<T*>(this);
		if (pT->GetState() & UAS_BLOCKED)
		{
			return S_OK;
		}

		if (!pPUU)
			return E_POINTER;

		if (!m_pOpenParentUndoUnit)
		{
			m_pOpenParentUndoUnit = pPUU;
			return S_OK;
		}

		return m_pOpenParentUndoUnit->Open(pPUU);
	}
	STDMETHOD(Close)(IOleParentUndoUnit* pPUU, BOOL fCommit)
	{
		T* pT = static_cast<T*>(this);
		if (pT->GetState() & UAS_BLOCKED)
		{
			CComPtr<IOleParentUndoUnit> spPUU;
			if ( spPUU.IsEqualObject(this) )
				return S_FALSE;
			return S_OK;
		}

		if (!m_pOpenParentUndoUnit) // No open parent unit, return S_FALSE
			return S_FALSE;

		HRESULT hRet = m_pOpenParentUndoUnit->Close(pPUU, fCommit);
		if (hRet == S_OK) // Handled by the child parent undo unit
			return S_OK;

		if (hRet == S_FALSE) // Verify that unit is equal to open child
		{
			if (! m_pOpenParentUndoUnit.IsEqualObject(pPUU) )
			{
				return E_INVALIDARG;
			}

			if (fCommit) // Add to the collection
			{
				m_childUnits.push_front(CAdapt<CComPtr<IOleUndoUnit> >(pPUU));
			}

			m_pOpenParentUndoUnit.Release();
			return S_OK;
		}

		return hRet;
	}
	STDMETHOD(Add)(IOleUndoUnit* pUU)
	{
		T* pT = static_cast<T*>(this);
		if (pT->GetState() & UAS_BLOCKED)
			return S_OK;
		
		if (!m_pOpenParentUndoUnit)
		{
			if (!m_childUnits.empty())
			{
				(*m_childUnits.begin()).m_T->OnNextAdd();
			}
			m_childUnits.push_front(CAdapt<CComPtr<IOleUndoUnit> >(pUU));
			return S_OK;
		}

		return m_pOpenParentUndoUnit->Add(pUU);
	}
	STDMETHOD(FindUnit)(IOleUndoUnit* pUU)
	{
		OleUndoUnitList::iterator iter = 
			std::find_if(m_childUnits.begin(), m_childUnits.end(), FindUnitFunction(pUU));

		return iter != m_childUnits.end() ? S_OK : S_FALSE;
	}
	STDMETHOD(GetParentState)(DWORD* pdwState)
	{
		if (!m_pOpenParentUndoUnit)
		{
			T* pT = static_cast<T*>(this);
			*pdwState = pT->GetState();
			return S_OK;
		}
		return m_pOpenParentUndoUnit->GetParentState(pdwState);
	}
};

/*
* Helper classes to create enumerators
*/
class _CopyIOleUndoUnitFromAdapt
{
public:
	static HRESULT copy(IOleUndoUnit** p1, CAdapt<CComPtr<IOleUndoUnit> >* p2)
	{
		*p1 = p2->m_T;
		if (*p1)
			(*p1)->AddRef();
		return S_OK;
	}
	static void init(IOleUndoUnit** ) {}
	static void destroy(IOleUndoUnit** p) {if (*p) (*p)->Release();}
};


typedef CComEnumOnSTL<IEnumOleUndoUnits, &IID_IEnumOleUndoUnits, IOleUndoUnit*,
			_CopyIOleUndoUnitFromAdapt, OleUndoUnitList > OleUndoUnitEnum;


/*
* Basic IOleUndoManager implementation
*
* TODO: Implement rollback when error occurs on UndoTo and RedoTo
*/
enum OleUndoManagerState
{
	oumBaseState,
	oumUndoState,
	oumRedoState
};

template<class T>
class ATL_NO_VTABLE COleUndoManagerImpl : public IOleUndoManager
{
public:
	COleUndoManagerImpl()
	{
		m_bEnabled = TRUE;
		m_mgrState = oumBaseState;
	}

	BOOL						m_bEnabled;
	CComPtr<IOleParentUndoUnit> m_pOpenParentUndoUnit;
	OleUndoUnitList				m_undoList;
	OleUndoUnitList				m_redoList;
	OleUndoManagerState			m_mgrState;

	inline void COleUndoManagerImpl_Add(IOleUndoUnit* pUU)
	{
		/* From the Platform SDK documentation.
		If the undo manager receives a new undo unit while in the base state, 
		it places the unit on the undo stack and discards the entire redo stack. 
		While it is in the undo state, it puts incoming units on the redo stack. 
		While it is in the redo state, it places them on the undo stack without 
		flushing the redo stack.
		*/

		switch (m_mgrState)
		{
			case oumBaseState:
				m_redoList.clear();
			case oumRedoState:
				m_undoList.push_front(CAdapt<CComPtr<IOleUndoUnit> >(pUU));
				break;
			case oumUndoState:
				m_redoList.push_front(CAdapt<CComPtr<IOleUndoUnit> >(pUU));
				break;
		}
	}

	STDMETHOD(Open)(IOleParentUndoUnit* pPUU)
	{
		if (!m_bEnabled) // If the UndoManager is disabled return S_OK immediately
			return S_OK;

		if (!pPUU)
			return E_POINTER;

		if (!m_pOpenParentUndoUnit)
		{
			m_pOpenParentUndoUnit = pPUU;
			return S_OK;
		}

		return m_pOpenParentUndoUnit->Open(pPUU);
	}
	STDMETHOD(Close)(IOleParentUndoUnit* pPUU, BOOL fCommit)
	{
		if (!m_bEnabled) // If the UndoManager is disabled return S_OK immediately
			return S_OK;

		if (!m_pOpenParentUndoUnit) // No open parent unit, return S_FALSE
			return S_FALSE;

		HRESULT hRet = m_pOpenParentUndoUnit->Close(pPUU, fCommit);
		if (hRet == S_OK) // Handled by the child parent undo unit
			return S_OK;

		if (hRet == S_FALSE) // Verify that unit is equal to open child
		{
			if (! m_pOpenParentUndoUnit.IsEqualObject(pPUU) )
			{
				return E_INVALIDARG;
			}

			if (fCommit) // Add to the collection
			{
				COleUndoManagerImpl_Add(pPUU);
			}

			m_pOpenParentUndoUnit.Release();
			return S_OK;
		}

		return hRet;
	}
	STDMETHOD(Add)(IOleUndoUnit* pUU)
	{
		if (!m_bEnabled)
			return S_OK;
		
		if (!m_pOpenParentUndoUnit)
		{
			// Notify last undo unit of new addition
			if ( m_mgrState == oumUndoState )
			{
				if (!m_redoList.empty())
				{
					(*m_redoList.begin()).m_T->OnNextAdd();
				}
			}
			else
			{
				if (!m_undoList.empty())
				{
					(*m_undoList.begin()).m_T->OnNextAdd();
				}
			}
			COleUndoManagerImpl_Add(pUU);
			return S_OK;
		}

		return m_pOpenParentUndoUnit->Add(pUU);
	}
	STDMETHOD(GetOpenParentState)(DWORD* pdwState)
	{
		if (!m_bEnabled)
		{
			*pdwState = UAS_BLOCKED;
			return S_OK;
		}
		
		if (!m_pOpenParentUndoUnit)
		{
			return S_FALSE;
		}

		return m_pOpenParentUndoUnit->GetParentState(pdwState);
	}
	STDMETHOD(DiscardFrom)(IOleUndoUnit* pUU)
	{
		if (!m_bEnabled)
			return E_UNEXPECTED;

		if (pUU == NULL && m_pOpenParentUndoUnit != NULL)
			m_pOpenParentUndoUnit.Release(); // Discard open unit

		OleUndoUnitList::iterator iter;

		if (pUU == NULL)
			iter = m_undoList.begin();
		else
			iter = std::find_if(m_undoList.begin(), m_undoList.end(), FindUnitFunction(pUU));

		if (iter != m_undoList.end() )
		{
			m_undoList.erase(iter, m_undoList.end());
			return S_OK;
		}

		if (pUU == NULL)
			iter = m_redoList.begin();
		else
			iter = std::find_if(m_redoList.begin(), m_redoList.end(), FindUnitFunction(pUU));

		if (iter != m_redoList.end() )
		{
			m_redoList.erase(iter, m_redoList.end());
			return S_OK;
		}

		return E_INVALIDARG;
	}
	STDMETHOD(UndoTo)(IOleUndoUnit* pUU)
	{
		if (!m_bEnabled)
			return E_UNEXPECTED;

		HRESULT hr = E_INVALIDARG;

		// Set manager to UndoState
		m_mgrState = oumUndoState;

		if (pUU == NULL)
		{
			if (!m_undoList.empty())
			{
				// Call Do on last Undo unit
				hr = (*m_undoList.begin()).m_T->Do(this);
				m_undoList.pop_front();
			}
		}
		else
		{
			// Call Do on Undo units
			OleUndoUnitList::iterator iter = 
				std::find_if(m_undoList.begin(), m_undoList.end(), FindUnitFunction(pUU));

			if (iter != m_undoList.end() )
			{
				// Iterator is incremented one
				hr = std::for_each(m_undoList.begin(), ++iter, CallDoFunction(this)).GetLastError();
				m_undoList.erase(m_undoList.begin(), iter);
			}
		}

		if (hr == E_ABORT)
		{
			// TODO: Rollback the Undo actions, clear both stacks and return E_FAIL
			hr = E_FAIL;
		}

		// Restore BaseState
		m_mgrState = oumBaseState;

		return hr;
	}
	STDMETHOD(RedoTo)(IOleUndoUnit* pUU)
	{
		if (!m_bEnabled)
			return E_UNEXPECTED;

		HRESULT hr = E_INVALIDARG;

		// Set manager to RedoState
		m_mgrState = oumRedoState;

		if (pUU == NULL)
		{
			if (!m_redoList.empty())
			{
				// Call Do on last Redo unit
				hr = (*m_redoList.begin()).m_T->Do(this);
				m_redoList.pop_front();
			}
		}
		else
		{
			// Call Do on Redo units
			OleUndoUnitList::iterator iter = 
				std::find_if(m_redoList.begin(), m_redoList.end(), FindUnitFunction(pUU));

			if (iter != m_redoList.end() )
			{
				// Iterator is incremented one
				hr = std::for_each(m_redoList.begin(), ++iter, CallDoFunction(this)).GetLastError();
				m_redoList.erase(m_redoList.begin(), iter);
			}
		}

		if (hr == E_ABORT)
		{
			// TODO: Rollback the Redo actions, clear both stacks and return E_FAIL
			hr = E_FAIL;
		}

		// Restore BaseState
		m_mgrState = oumBaseState;

		return hr;
	}
	STDMETHOD(EnumUndoable)(IEnumOleUndoUnits** ppEnum)
	{
		if (!ppEnum)
			return E_POINTER;
		*ppEnum = NULL;

		CComObject<OleUndoUnitEnum>* pEnum = NULL;
		HRESULT hr = pEnum->CreateInstance(&pEnum);
		
		if (SUCCEEDED(hr))
		{
			pEnum->AddRef();

			hr = pEnum->Init(this, m_undoList);
			
			if (SUCCEEDED(hr))
				hr = pEnum->QueryInterface(ppEnum);

			pEnum->Release();
		}

        return hr;
	}
	STDMETHOD(EnumRedoable)(IEnumOleUndoUnits** ppEnum)
	{
		if (!ppEnum)
			return E_POINTER;
		*ppEnum = NULL;

		CComObject<OleUndoUnitEnum>* pEnum = NULL;
		HRESULT hr = pEnum->CreateInstance(&pEnum);
		
		if (SUCCEEDED(hr))
		{
			pEnum->AddRef();

			hr = pEnum->Init(this, m_redoList);

			if (SUCCEEDED(hr))
				hr = pEnum->QueryInterface(ppEnum);

			pEnum->Release();
		}

        return hr;
	}
	STDMETHOD(GetLastUndoDescription)(BSTR* pBstr)
	{
		if (!m_bEnabled)
		{
			*pBstr = NULL;
			return E_UNEXPECTED;
		}

		if (!m_undoList.empty())
		{
			return (*m_undoList.begin()).m_T->GetDescription(pBstr);
		}

		return E_FAIL;
	}
	STDMETHOD(GetLastRedoDescription)(BSTR* pBstr)
	{
		if (!m_bEnabled)
		{
			*pBstr = NULL;
			return E_UNEXPECTED;
		}

		if (!m_redoList.empty())
		{
			return (*m_redoList.begin()).m_T->GetDescription(pBstr);
		}

		return E_FAIL;
	}
	STDMETHOD(Enable)(BOOL fEnable)
	{
		m_bEnabled = fEnable;
		return S_OK;
	}
};

}; // End namespace

#ifndef ATLOLE_NO_USING_NAMESPACE
using namespace ATLOLE;
#endif

#endif UNDOMGR_H__831D7338_E2BC_4405_BBF2_9FBB761AE1E9