Tree Control with Bitmap Checkboxes Supported

// XMLTree.cpp : implementation file
//

#include "stdafx.h"
#include "vs_treectrl.h"
#include <comdef.h>
#include "XMLTree.h"

#ifdef _DEBUG
#define new DEBUG_NEW
#undef THIS_FILE
static char THIS_FILE[] = __FILE__;
#endif

/////////////////////////////////////////////////////////////////////////////
// CXMLTree

CXMLTree::CXMLTree()
{
	m_pXML=NULL;
}

CXMLTree::~CXMLTree()
{
}


BEGIN_MESSAGE_MAP(CXMLTree, CTreeCtrl)
	//{{AFX_MSG_MAP(CXMLTree)
		ON_NOTIFY_REFLECT(TVN_ITEMEXPANDING, OnItemexpanding)
	ON_WM_LBUTTONDOWN()
	ON_WM_KEYDOWN()
	//}}AFX_MSG_MAP
END_MESSAGE_MAP()

/////////////////////////////////////////////////////////////////////////////
// CXMLTree message handlers
void CXMLTree::OnItemexpanding(NMHDR* pNMHDR, LRESULT* pResult) 
{
	NM_TREEVIEW* pNMTreeView = (NM_TREEVIEW*)pNMHDR;
	HTREEITEM hItem = pNMTreeView->itemNew.hItem;
	/*MSXML::*/IXMLDOMElement* node = (/*MSXML::*/IXMLDOMElement*)/*GetTreeCtrl().*/GetItemData(hItem);
	HRESULT hr;

	*pResult = 0;
	
	// If m_bOptimizeMemory is false then we optimize by not adding and deleting 
	// already added items. The trade-off is memory exhausted!
	
	CWaitCursor waitCursor;     // This could take a while!
	/*GetTreeCtrl().*/LockWindowUpdate();
	if (pNMTreeView->action == TVE_EXPAND) {
		if (m_bOptimizeMemory == FALSE) {
			HTREEITEM hChildItem;
			if ((hChildItem = /*GetTreeCtrl().*/GetChildItem(hItem)) != NULL) {
				/*MSXML::*/IXMLDOMElement* childNode = 
					(/*MSXML::*/IXMLDOMElement*)/*GetTreeCtrl().*/GetItemData(hChildItem);
				if (childNode == NULL) {
					/*GetTreeCtrl().*/DeleteItem(hChildItem);
					/*MSXML::*/IXMLDOMNode* firstChild = NULL;
					hr = node->get_firstChild(&firstChild);
					if (SUCCEEDED(hr) && firstChild != NULL) {
						if (populateNode((/*MSXML::*/IXMLDOMElement*)firstChild, hItem) == FALSE) {
							*pResult = 1;            
						}
					}
				}
			}
		} else {
			deleteFirstChild(hItem);
			/*MSXML::*/IXMLDOMNode* firstChild = NULL;
			hr = node->get_firstChild(&firstChild);
			if (SUCCEEDED(hr) && firstChild != NULL) {
				if (populateNode((/*MSXML::*/IXMLDOMElement*)firstChild, hItem) == FALSE) {
					*pResult = 1;            
				}
			}
		}
	} else { // pNMTreeView->action == TVE_COLLAPSE
		if (m_bOptimizeMemory == TRUE) {
			deleteAllChildren(hItem);
			// Set button state.
			VARIANT_BOOL vtb;
			//if (node->hasChildNodes()) {
			node->hasChildNodes (&vtb);
			if(vtb)
			{
				int nImage, nSelectedImage;
				nImage = nSelectedImage = getIconIndex(node);
				HTREEITEM hChildItem = /*GetTreeCtrl().*/InsertItem(_T(""), nImage, 
					nSelectedImage, hItem);
				/*GetTreeCtrl().*/SetItemData(hChildItem, (DWORD)NULL);
			}
		}
	}
	/*GetTreeCtrl().*/UnlockWindowUpdate();
}


BOOL CXMLTree::loadXML(const CString &strPathName, const BOOL bOptimizeMemory /*= FALSE*/ ,const BOOL need_doc)
{
	/*MSXML::*/IXMLDOMDocument* document = NULL;
	/*MSXML::*/IXMLDOMParseError* parseError = NULL;
	/*MSXML::*/IXMLDOMElement* element = NULL;	

	DeleteAllItems();

	HRESULT hr;
	hr = CoInitialize(NULL);
	if (FAILED(hr)) {
		return FALSE;
	}

	hr = CoCreateInstance(/*MSXML::*/CLSID_DOMDocument, NULL, 
		CLSCTX_INPROC_SERVER | CLSCTX_LOCAL_SERVER, 
		/*MSXML::*/IID_IXMLDOMDocument, (LPVOID*)&document);
	if (!document) {
		return FALSE;
	}

	BSTR bstr = NULL;
	document->put_async(VARIANT_FALSE);
	bstr = strPathName.AllocSysString();
	VARIANT_BOOL varOkay;
	//= document->load(bstr);
	document->load (_variant_t(bstr),&varOkay);
	SysFreeString(bstr);
	
	BSTR nodeName;
	HTREEITEM hItem;	
	if (varOkay) {
		hr = document->get_documentElement(&element);
		if (FAILED(hr) || element == NULL) {
			MessageBox(_T("Empty document!"), _T("Error Loading XML"), MB_ICONWARNING);
			return FALSE;
		}
		element->get_nodeName(&nodeName);
		hItem = insertItem(element, CString(nodeName), ILI_ROOT, ILI_ROOT);	
		// Optional expand the root initially.
		Expand(hItem, TVE_EXPAND);
		Select(hItem, TVGN_CARET);
	} else {
		long line, linePos;
		BSTR reason = NULL;
		
		document->get_parseError(&parseError);
		parseError->get_errorCode(&hr);
		parseError->get_line(&line);
		parseError->get_linepos(&linePos);
		parseError->get_reason(&reason);
		
		CString strMsg;
		strMsg.Format(_T("Error 0x%.8X on line %d, position %d\r\nReason: %s"), 
			hr, line, linePos, CString(reason));
		MessageBox(strMsg, _T("Error Loading XML"), MB_ICONWARNING);
				
		SysFreeString(reason);
		return FALSE;
	}

	// Must initialize here.
	m_bOptimizeMemory = bOptimizeMemory;
	if ( need_doc)
	{
		m_pXML = document;
		m_URL = strPathName;
	}
	return TRUE;
}

// This function populates all siblings of the [in] node.
//
BOOL CXMLTree::populateNode(/*MSXML::*/IXMLDOMElement *node, const HTREEITEM &hParent)
{
	HRESULT hr = S_OK;
	BSTR nodeType, nodeName;
	HTREEITEM hItem;

	node->get_nodeTypeString(&nodeType);
	if (!wcscmp(nodeType, L"element")) {
		node->get_nodeName(&nodeName);
		hItem = insertItem(node, CString(nodeName), ILI_ELEMENT, ILI_ELEMENT, hParent);
		populateAttributes(node, hItem);
	} else if(!wcscmp(nodeType, L"text")) {
		node->get_text(&nodeName);
		hItem = insertItem(node, CString(nodeName), ILI_TEXT, ILI_TEXT, hParent);	
	} else if(!wcscmp(nodeType, L"comment")) {
		node->get_nodeName(&nodeName);
		hItem = insertItem(node, CString(nodeName), ILI_COMMENT, ILI_COMMENT, hParent);	
	} else {    // Handle more data types here if needed.
		node->get_nodeName(&nodeName);
		hItem = insertItem(node, CString(nodeName), ILI_OTHER, ILI_OTHER, hParent);	
	}
	
	/*MSXML::*/IXMLDOMNode* nextSibling = NULL;
	hr = node->get_nextSibling(&nextSibling);
	if (SUCCEEDED(hr) && nextSibling != NULL) {
		populateNode((/*MSXML::*/IXMLDOMElement*)nextSibling, hParent);
	}

	return TRUE;
}

BOOL CXMLTree::populateAttributes(/*MSXML::*/IXMLDOMElement *node, const HTREEITEM &hParent)
{
	HRESULT hr = S_OK;
	BSTR nodeName;
	HTREEITEM hItem;

	/*MSXML::*/IXMLDOMNamedNodeMap* namedNodeMap = NULL;
	hr = node->get_attributes(&namedNodeMap);
	if (SUCCEEDED(hr) && namedNodeMap != NULL) {
		long listLength;
		hr = namedNodeMap->get_length(&listLength);
		for(long i = 0; i < listLength; i++) {
			/*MSXML::*/IXMLDOMNode* listItem = NULL;
			hr = namedNodeMap->get_item(i, &listItem);
			listItem->get_nodeName(&nodeName);
			hItem = insertItem((/*MSXML::*/IXMLDOMElement*)listItem, 
				CString(nodeName), ILI_ATTRIBUTE, ILI_ATTRIBUTE, hParent);	
		}
	}

	return TRUE;
}

HTREEITEM CXMLTree::insertItem(/*MSXML::*/IXMLDOMElement *node, 
	const CString &nodeName, int nImage, int nSelectedImage, 
	HTREEITEM hParent /*= TVI_ROOT*/, HTREEITEM hInsertAfter /*= TVI_LAST*/)
{
	HTREEITEM hItem = /*GetTreeCtrl().*/InsertItem(nodeName, nImage, 
		nSelectedImage, hParent, hInsertAfter);	


	if ( GetWindowLong(m_hWnd ,GWL_STYLE) & TVS_CHECKBOXES )
	{	
		BOOL checked=FALSE;
		if ( hParent != TVI_ROOT)
			 checked = IsItemChecked(hParent);
		if ( checked )
			SetChecked(hItem);
		else
			SetUnchecked(hItem);
	}
	/*GetTreeCtrl().*/SetItemData(hItem, (DWORD)node);
	

	// Set button state.
	VARIANT_BOOL vtb;
	node->hasChildNodes(&vtb);
	if (vtb) {
		HTREEITEM hChildItem = /*GetTreeCtrl().*/InsertItem(_T(""), nImage, 
			nSelectedImage, hItem);
		/*GetTreeCtrl().*/SetItemData(hChildItem, (DWORD)NULL);
	}

	return hItem;
}

void CXMLTree::deleteFirstChild(const HTREEITEM &hItem)
{
	HTREEITEM hChildItem;
	if ((hChildItem = /*GetTreeCtrl().*/GetChildItem(hItem)) != NULL) {
		/*GetTreeCtrl().*/DeleteItem(hChildItem);
	}
}

 void CXMLTree::deleteAllChildren(const HTREEITEM &hItem)
 {
 	HTREEITEM hChildItem;
 	if ((hChildItem = /*GetTreeCtrl().*/GetChildItem(hItem)) == NULL) {
 		return;
 	}
 
 	do {
 		HTREEITEM hNextItem = /*GetTreeCtrl().*/GetNextSiblingItem(hChildItem);
 		/*GetTreeCtrl().*/DeleteItem(hChildItem);
 		hChildItem = hNextItem;
 	} while (hChildItem != NULL);
 }

int CXMLTree::getIconIndex(/*MSXML::*/IXMLDOMElement *node)
{
	BSTR nodeType;

	node->get_nodeTypeString(&nodeType);
	if (!wcscmp(nodeType, L"element")) {
		return ILI_ELEMENT;
	} else if(!wcscmp(nodeType, L"attribute")) {
		return ILI_ATTRIBUTE;
	} else if(!wcscmp(nodeType, L"text")) {
		return ILI_TEXT;
	} else if(!wcscmp(nodeType, L"comment")) {
		return ILI_COMMENT;
	}

	return ILI_OTHER;
}



BOOL CXMLTree::Create(LPCTSTR lpszClassName, LPCTSTR lpszWindowName, DWORD dwStyle, const RECT& rect, CWnd* pParentWnd, UINT nID, CCreateContext* pContext) 
{
	// TODO: Add your specialized code here and/or call the base class
	
	return CWnd::Create(lpszClassName, lpszWindowName, dwStyle, rect, pParentWnd, nID, pContext);
}

BOOL CXMLTree::CreateTree(CRect rect,CWnd *parent, UINT image)
{
	if(!CTreeCtrl::Create (WS_TABSTOP | WS_CHILD | WS_BORDER|WS_THICKFRAME   
   | TVS_HASBUTTONS | TVS_LINESATROOT | TVS_HASLINES
   | TVS_DISABLEDRAGDROP
	,rect,parent,65535))
		return FALSE;
	if (image)
		SetImages (image);
	return TRUE;
}

void CXMLTree::SetImages(UINT resource)
{
	m_theImageList.Detach();	// For MDI
	if ( resource ==0 )
		m_theImageList.Create(IDB_BITMAP_TREE_ICONS, 16, 4, RGB(255, 255, 255));	
	else
		m_theImageList.Create (resource,16,4,RGB(255,255,255));
	SetImageList(&m_theImageList, TVSIL_NORMAL);

	m_imageState.Detach ();

	m_imageState.Create( IDB_STATE, 16, 1, RGB(255,255,255) );
	SetImageList( &m_imageState, TVSIL_STATE );


}

void CXMLTree::SaveXML(const CString &strFileName)
{
	CString fname = m_URL;
	if( !strFileName.IsEmpty ())
	{
		fname = strFileName;
	}
	m_pXML->save (_variant_t(fname));
}

IXMLDOMElement* CXMLTree::AddNewNode(CString nodeName, CString strValue)
{
	HTREEITEM item = GetSelectedItem ();
	/*MSXML::*/IXMLDOMElement *node=NULL;
	/*MSXML::*/IXMLDOMElement *new_node=NULL;
	if ( item )
	{
		node= (/*MSXML::*/IXMLDOMElement*)GetItemData (item);
		m_pXML->createElement (_bstr_t(nodeName),&new_node);
		if ( new_node)
		{
			//new_node->put_nodeValue (_variant_t(strValue));
			new_node->put_text (_bstr_t(strValue));
		}
		if ( node)
		{
			IXMLDOMNode *element=NULL;
			node->appendChild (new_node,&element);
			HTREEITEM new_item = insertItem(new_node,nodeName,1,1,item,TVI_LAST);
			SelectItem (new_item);
			return new_node;
		}

	}
	return NULL;
}

void CXMLTree::DeleteNode(CString strPath)
{
	HTREEITEM item = GetSelectedItem ();
	IXMLDOMElement *node=NULL;
	//if (strPath.IsEmpty ())
	if( item)
	{
		 node= (IXMLDOMElement*)GetItemData (item);
		if ( node )
		{
			IXMLDOMNode *parent=NULL;
			IXMLDOMNode *old_node=NULL;
			node->get_parentNode (&parent);		
			if ( parent)
				parent->removeChild (node,&old_node);
		}
		this->DeleteItem (item);
	}
}

void CXMLTree::AddAttribute(IXMLDOMElement *element,CString strName, CString strValue)
{
	IXMLDOMAttribute *attribute=NULL;
	m_pXML->createAttribute (_bstr_t(strName),&attribute);
	IXMLDOMAttribute *element1=NULL;
	element->setAttributeNode  (attribute,&element1);
	element->setAttribute (_bstr_t(strName),_variant_t(strValue));

}
void CXMLTree::AddXMLToNode(HTREEITEM item ,IXMLDOMElement *element,CString strXML)
{
	IXMLDOMNode *ele_new=NULL;
	IXMLDOMDocument *document=NULL;
	IXMLDOMElement *origin=NULL;
	HRESULT hr;
	if ( item )
	{

		BSTR bstr;
		bstr = strXML.AllocSysString ();
		hr = CoCreateInstance(CLSID_DOMDocument, NULL, 
		CLSCTX_INPROC_SERVER | CLSCTX_LOCAL_SERVER, 
		IID_IXMLDOMDocument, (LPVOID*)&document);
		if (!document) 	return ;
		document->put_async(VARIANT_FALSE);
		VARIANT_BOOL varOkay;
		document->loadXML (bstr,&varOkay);
		SysFreeString(bstr);
		if (varOkay) 
		{
			hr = document->get_documentElement(&origin);
			if (SUCCEEDED(hr))
			{
				hr=origin->cloneNode (VARIANT_TRUE,&ele_new);
				if (SUCCEEDED(hr))
				{
					IXMLDOMNode *old_node=NULL;
					hr= element->appendChild (ele_new,&old_node);
					if (SUCCEEDED(hr))
					{
					   BSTR name;
					   ele_new->get_nodeName (&name);
					   HTREEITEM new_item =insertItem((IXMLDOMElement*)ele_new,CString(name),0,0,item,TVI_LAST);
					   SelectItem(new_item);
					}


				}

			}
		}
}

}

	

void CXMLTree::OnLButtonDown(UINT nFlags, CPoint point) 
{
	UINT uFlags=0;
	HTREEITEM hti = HitTest(point,&uFlags);

	if( uFlags & TVHT_ONITEMSTATEICON )
	{
		ToggleItemState(hti);
//		int iImage = GetItemState( hti, TVIS_STATEIMAGEMASK )>>12;
//		SetItemState( hti, INDEXTOSTATEIMAGEMASK(iImage == 1 ? 2 : 1), 
//					TVIS_STATEIMAGEMASK );
		return;
	}
	
	CTreeCtrl::OnLButtonDown(nFlags, point);
}

void CXMLTree::OnKeyDown(UINT nChar, UINT nRepCnt, UINT nFlags) 
{
	if( nChar == VK_SPACE )
	{
		HTREEITEM hti = GetSelectedItem();
		ToggleItemState(hti);
		//int iImage = GetItemState( hti, TVIS_STATEIMAGEMASK )>>12;
		//SetItemState( hti, INDEXTOSTATEIMAGEMASK(iImage == 1 ? 2 : 1),
		//			TVIS_STATEIMAGEMASK );
		return;
	}
	CTreeCtrl::OnKeyDown(nChar, nRepCnt, nFlags);
}

BOOL CXMLTree::IsItemChecked(HTREEITEM hItem)
{
	return GetItemState( hItem, TVIS_STATEIMAGEMASK )>>12 == 2;
}



HTREEITEM CXMLTree::GetFirstCheckedItem()
{
	for ( HTREEITEM hItem = GetRootItem(); hItem!=NULL; hItem = GetNextItem( hItem ) )
		if ( IsItemChecked(hItem) )
			return hItem;

	return NULL;

}

HTREEITEM CXMLTree::GetNextCheckedItem(HTREEITEM hItem)
{
	for ( hItem = GetNextItem( hItem ); hItem!=NULL; hItem = GetNextItem( hItem ) )
		if ( IsItemChecked(hItem) )
			return hItem;

	return NULL;

}

HTREEITEM CXMLTree::GetPrevCheckedItem(HTREEITEM hItem)
{
	for ( hItem = GetPrevItem( hItem ); hItem!=NULL; hItem = GetPrevItem( hItem ) )
		if ( IsItemChecked(hItem) )
			return hItem;

	return NULL;

}

void CXMLTree::SetChecked(HTREEITEM hItem)
{
	SetItemState( hItem, INDEXTOSTATEIMAGEMASK(2), TVIS_STATEIMAGEMASK );

}
// GetNextItem	- Get next item as if outline was completely expanded
// Returns		- The item immediately below the reference item
// hItem		- The reference item

HTREEITEM CXMLTree::GetNextItem(HTREEITEM hItem)
{

	HTREEITEM	hti;

	if( ItemHasChildren( hItem ) )
		return GetChildItem( hItem );		// return first child
	else{
		// return next sibling item
		// Go up the tree to find a parent's sibling if needed.
		while( (hti = GetNextSiblingItem( hItem )) == NULL ){
			if( (hItem = GetParentItem( hItem ) ) == NULL )
				return NULL;
		}
	}
	return hti;


}

HTREEITEM CXMLTree::GetPrevItem(HTREEITEM hItem)
{
	HTREEITEM hti;

	// Return a previous sibling item if it exists
	if( hti = GetPrevSiblingItem(hItem) )
		return hti;		
	else
		// No sibling, return the parent 
		return GetParentItem(hItem);

}

void CXMLTree::ToggleItemState(HTREEITEM hti)
{
	int iImage = GetItemState( hti, TVIS_STATEIMAGEMASK )>>12;
	SetItemState( hti, INDEXTOSTATEIMAGEMASK(iImage == 1 ? 2 : 1), 
				TVIS_STATEIMAGEMASK );
	if ( ItemHasChildren(hti))
	{
		HTREEITEM htiChild = this->GetChildItem (hti);
		if (htiChild)
			ToggleItemState(htiChild);
		else
			return ;
		HTREEITEM htiSibling = GetNextSiblingItem (htiChild);
		while (htiSibling )
		{
			ToggleItemState(htiSibling);
			htiSibling = GetNextSiblingItem(htiSibling);
		}


	}

}

void CXMLTree::SetUnchecked(HTREEITEM hti)
{
	SetItemState( hti, INDEXTOSTATEIMAGEMASK(1), TVIS_STATEIMAGEMASK );

}