Custom Drawn Vertical Tree Control

/**
 * Copyright (C) 2005
 * Doga Arinir
 * 
 * Author: Doga Arinir
 * E-Mail: arinir@gmx.de
 *
 * This software is provided 'as-is', without any express or implied
 * warranty.  In no event will the author or the company be held liable 
 * for any damages arising from the use of this software. EXPECT BUGS!
 * 
 * You may use this software in compiled form in any way you desire PROVIDING it is
 * not sold for profit without the authors written permission, and providing that this
 * notice and the authors name is included. If the source code in this file is used in 
 * any commercial application then acknowledgement must be made to the author of this file.
 */

#include "stdafx.h"
#include "VerticalTree.h"
#include "CustomVerticalTree.h"

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

/////////////////////////////////////////////////////////////////////////////
// CVerticalTree

CVerticalTree::CVerticalTree() : CTreeCtrl(), m_bCustomDrawn(TRUE), m_ItemSpaceX(10), m_ItemSpaceY(250)
{
	m_TreeLinePen.CreatePen(PS_DOT, 1, RGB(120,120,120));

	m_LeafTextColor = RGB(0,0,0);
	m_LeafBkColor = RGB(240, 240, 240);

	m_ItemTextColor = RGB(255, 255, 0);
	m_ItemBkColor = RGB(200, 200, 255);
}

CVerticalTree::~CVerticalTree()
{
	DeleteAllViewportItems();
}
void CVerticalTree::SetCustomDrawn(BOOL c)
{
	m_bCustomDrawn = c;
}
BOOL CVerticalTree::GetCustomDrawn()
{
	return m_bCustomDrawn;
}


BEGIN_MESSAGE_MAP(CVerticalTree, CTreeCtrl)
	//{{AFX_MSG_MAP(CVerticalTree)
	ON_WM_PAINT()
	ON_WM_ERASEBKGND()
	ON_WM_HSCROLL()
	ON_WM_LBUTTONUP()
	ON_WM_MOUSEWHEEL()
	ON_WM_VSCROLL()
	ON_WM_LBUTTONDOWN()
	ON_WM_MOUSEMOVE()
	ON_WM_SIZING()
	ON_WM_SIZE()
	//}}AFX_MSG_MAP
END_MESSAGE_MAP()

/////////////////////////////////////////////////////////////////////////////
// Behandlungsroutinen f�r Nachrichten CVerticalTree 

void CVerticalTree::OnPaint() 
{
	if (m_bCustomDrawn)
	{
		CPaintDC dc(this);
		//Double Buffering...
		CDC dc_backbuffer;
		CBitmap bb_image;

		//Calculate scoll position and control dimensions
		UpdateDimensions();

		//Create device context for the backbuffer
		dc_backbuffer.CreateCompatibleDC(&dc);
		//Create the image we are going to draw in
		bb_image.CreateCompatibleBitmap(&dc, m_Rect.Width(), m_Rect.Height());
		dc_backbuffer.SelectObject(&bb_image);

		//We use the same Font which was configured in the main device context
		dc_backbuffer.SelectObject(GetFont());
		dc_backbuffer.SetBkMode(TRANSPARENT);

		//General Info
		DWORD style = GetStyle();
		m_bHasButtons = (style & TVS_HASBUTTONS) != 0;
		m_bHasLines = (style & TVS_HASLINES) != 0;

		//Now draw the items...
		m_XMax = 0;
		m_YMax = 0;
		dc_backbuffer.FillSolidRect(m_Rect, dc_backbuffer.GetBkColor()); 
		int xoffset = 0;
		for (HTREEITEM root = GetRootItem(); root != NULL; root = GetNextItem(root, TVGN_NEXT))
			xoffset = DrawItem(&dc_backbuffer, root, xoffset);

		//Copy the backbuffer to the main device context...
		dc.BitBlt( m_Rect.left, m_Rect.top, m_Rect.Width(), m_Rect.Height(), &dc_backbuffer, 0, 0, SRCCOPY);

		ResetScrollbars();
	}
	else
		CTreeCtrl::OnPaint();
}

BOOL CVerticalTree::OnEraseBkgnd(CDC* pDC) 
{
	if (m_bCustomDrawn)
		return TRUE;
	else
		return CTreeCtrl::OnEraseBkgnd(pDC);
}

void CVerticalTree::UpdateDimensions()
{
	GetClientRect(&m_Rect);
	SCROLLINFO scroll_info;
	if (GetScrollInfo(SB_HORZ, &scroll_info, SIF_POS | SIF_RANGE))
	{
		m_OffsetX = -scroll_info.nPos;
	}
	else
	{
		m_OffsetX = m_Rect.left;
	}
	if (GetScrollInfo(SB_VERT, &scroll_info, SIF_POS | SIF_RANGE))
	{
		if ( scroll_info.nMin == 0 && scroll_info.nMax == 100) 
			scroll_info.nMax = 0;

		m_OffsetY = -scroll_info.nPos;
	}
	else
	{
		m_OffsetY = m_Rect.top;
	}
}


void CVerticalTree::DrawItemText(CDC *pDC, HTREEITEM item, const char *text, int x, int y, int cx, int cy, COLORREF bkcolor,COLORREF txtcolor)
{
	pDC->FillSolidRect(x+m_OffsetX, y+m_OffsetY, cx, cy, bkcolor);
	pDC->SetTextColor(txtcolor);
	pDC->TextOut(x+m_OffsetX,y+m_OffsetY,text);


	ItemViewport *nvp = GetViewport(item);
	
	nvp->x = x;
	nvp->y = y;
	nvp->cx = cx;
	nvp->cy = cy;

	if (x + cx > m_XMax)
		m_XMax = x + cx;
	if (y + cy > m_YMax)
		m_YMax = y + cy;

}


int CVerticalTree::DrawItem(CDC *pDC, HTREEITEM item, int x, int level)
{
	CString name = GetItemText(item);
	CSize text_size = pDC->GetTextExtent(name);

	int state = GetItemState(item, TVIF_STATE);
	bool selected = (state & TVIS_SELECTED) != 0;

	if (ItemHasChildren(item))
	{
		int left = x;
		int right = 0;
		int childcount = 0;

		if (state & TVIS_EXPANDED)
			for (HTREEITEM childitem = GetChildItem(item); childitem != NULL; childitem = GetNextItem(childitem, TVGN_NEXT))
			{
				right = DrawItem(pDC, childitem, x, level + 1);
				x = right;
				childcount++;
			}

		right = right - m_ItemSpaceX;
		int width =  right - left;
		x = left + width / 2 - text_size.cx / 2;

		if (x < left + m_ItemSpaceX)
		   x = left;
		

		int y = level*m_ItemSpaceY;
		DrawItemText(pDC, item, name, x, level*m_ItemSpaceY, text_size.cx, text_size.cy, m_ItemBkColor, m_ItemTextColor);

		//Draw lines...
		if (m_bHasLines && (state & TVIS_EXPANDED))
		{
			int xstart = x + text_size.cx / 2;
			int ystart = y + text_size.cy + BUTTON_SIZE;
			CGdiObject *oldPen = pDC->SelectObject(&m_TreeLinePen);
			for (HTREEITEM childitem = GetChildItem(item); childitem != NULL; childitem = GetNextItem(childitem, TVGN_NEXT))
			{
				ItemViewport *current = GetViewport(childitem);
				pDC->MoveTo(xstart+m_OffsetX,ystart+m_OffsetY);
				pDC->LineTo(current->x +m_OffsetX + current->cx / 2 , current->y + m_OffsetY);
			}
			pDC->SelectObject(oldPen);
		}

		if (m_bHasButtons)
		{
			pDC->Draw3dRect(x+m_OffsetX+text_size.cx / 2 - BUTTON_SIZE / 2, y+m_OffsetY + text_size.cy, BUTTON_SIZE, BUTTON_SIZE, RGB(200,200,200), RGB(100,100,100));
			pDC->MoveTo(x+m_OffsetX+text_size.cx / 2 - BUTTON_SIZE / 2 + 2, y+m_OffsetY + text_size.cy + BUTTON_SIZE / 2);
			pDC->LineTo(x+m_OffsetX+text_size.cx / 2 + BUTTON_SIZE / 2 - 1, y+m_OffsetY + text_size.cy + BUTTON_SIZE / 2);

			if ((state & TVIS_EXPANDED) == 0)
			{
				pDC->MoveTo(x+m_OffsetX+text_size.cx / 2, y+m_OffsetY + text_size.cy + 2);
				pDC->LineTo(x+m_OffsetX+text_size.cx / 2, y+m_OffsetY + text_size.cy + BUTTON_SIZE - 2);
			}
		}
		if (right > x + text_size.cx)
		   return right + m_ItemSpaceX;
		else
		   return x + text_size.cx + m_ItemSpaceX;
	}
	else
	{
		DrawItemText(pDC, item, name, x, level*m_ItemSpaceY, text_size.cx, text_size.cy, m_LeafBkColor, m_LeafTextColor);
		return x + text_size.cx + m_ItemSpaceX;
	}
}

void CVerticalTree::SetItemSpaceX(int space)
{
	m_ItemSpaceX = space;
}

int CVerticalTree::GetItemSpaceX()
{
	return m_ItemSpaceX;
}
void CVerticalTree::SetItemSpaceY(int space)
{
	m_ItemSpaceY = space;
}

int CVerticalTree::GetItemSpaceY()
{
	return m_ItemSpaceY;
}

void CVerticalTree::OnHScroll(UINT nSBCode, UINT nPos, CScrollBar* pScrollBar) 
{
	if (m_bCustomDrawn)
	{
		int iScrollBarPos = GetScrollPos( SB_HORZ );

		switch( nSBCode )
		{
			case SB_LINEUP:
				iScrollBarPos = max( iScrollBarPos - 1, 0 );
			break;

			case SB_LINEDOWN:
				iScrollBarPos = min( iScrollBarPos + 1, GetScrollLimit( SB_HORZ ) );
			break;

			case SB_PAGEUP:
				iScrollBarPos = max( iScrollBarPos - m_Rect.Width(), 0 );
			break;

			case SB_PAGEDOWN:
				iScrollBarPos = min( iScrollBarPos + m_Rect.Width(), GetScrollLimit( SB_HORZ ) );
			break;

			case SB_THUMBTRACK:
			case SB_THUMBPOSITION:
				iScrollBarPos = nPos;
				break;
		}		

		SetScrollPos( SB_HORZ, iScrollBarPos );
		
		Invalidate();
	}
	else
		CTreeCtrl::OnHScroll(nSBCode, nPos, pScrollBar);
}

void CVerticalTree::OnLButtonUp(UINT nFlags, CPoint point) 
{
	if (m_bCustomDrawn)
	{
		HTREE2VPMAP::iterator it = m_HTree2VP.begin();

		point.x -= m_OffsetX;
		point.y -= m_OffsetY;
		while (it != m_HTree2VP.end())
		{
			ItemViewport *vp = (*it).second;
			if (vp->x <= point.x && ((vp->x+vp->cx) >= point.x)
				&& vp->y <= point.y && ((vp->y+vp->cy) >= point.y))
			{
				Expand(vp->item, TVE_TOGGLE);
				Invalidate();
				return;
			}
			it++;
		}
	}
	else
		CTreeCtrl::OnLButtonUp(nFlags, point);
}

BOOL CVerticalTree::OnMouseWheel(UINT nFlags, short zDelta, CPoint pt) 
{
	if (m_bCustomDrawn)
	{
		return FALSE;
	}
	else	
		return CTreeCtrl::OnMouseWheel(nFlags, zDelta, pt);
}

void CVerticalTree::OnVScroll(UINT nSBCode, UINT nPos, CScrollBar* pScrollBar) 
{
	if (m_bCustomDrawn)
	{
		int iScrollBarPos = GetScrollPos( SB_VERT );

		switch( nSBCode )
		{
			case SB_LINEUP:
				iScrollBarPos = max( iScrollBarPos - 1, 0 );
			break;

			case SB_LINEDOWN:
				iScrollBarPos = min( iScrollBarPos + 1, GetScrollLimit( SB_VERT ) );
			break;

			case SB_PAGEUP:
				iScrollBarPos = max( iScrollBarPos - m_Rect.Height(), 0 );
			break;

			case SB_PAGEDOWN:
				iScrollBarPos = min( iScrollBarPos + m_Rect.Height(), GetScrollLimit( SB_VERT ) );
			break;

			case SB_THUMBTRACK:
			case SB_THUMBPOSITION:
				iScrollBarPos = nPos;
				break;
		}		

		SetScrollPos( SB_VERT, iScrollBarPos );
		
		Invalidate();
	}
	else
		CTreeCtrl::OnVScroll(nSBCode, nPos, pScrollBar);
}




void CVerticalTree::OnLButtonDown(UINT nFlags, CPoint point) 
{
	if (m_bCustomDrawn)
	{
	}
	else	
		CTreeCtrl::OnLButtonDown(nFlags, point);
}

void CVerticalTree::OnMouseMove(UINT nFlags, CPoint point) 
{
	if (m_bCustomDrawn)
	{
	}
	else	
		CTreeCtrl::OnMouseMove(nFlags, point);
}


HTREEITEM CVerticalTree::InsertItem(LPTVINSERTSTRUCT lpInsertStruct)
{
	HTREEITEM htreeitem = CTreeCtrl::InsertItem(lpInsertStruct);
	//I presume that CTreeCtrl::InsertItem won't create same handles,
	//if the key is in use this will lead to memory leaks...
	m_HTree2VP.insert(HTREE2VPMAP::value_type(htreeitem, new ItemViewport(htreeitem)));
	return htreeitem;
}
HTREEITEM CVerticalTree::InsertItem(UINT nMask, LPCTSTR lpszItem, int nImage, int nSelectedImage, UINT nState, UINT nStateMask, LPARAM lParam, HTREEITEM hParent, HTREEITEM hInsertAfter)
{
	HTREEITEM htreeitem = CTreeCtrl::InsertItem(nMask, lpszItem, nImage, nSelectedImage, nState, nStateMask, lParam, hParent, hInsertAfter);
	//I presume that CTreeCtrl::InsertItem won't create same handles,
	//if the key is in use this will lead to memory leaks...
	m_HTree2VP.insert(HTREE2VPMAP::value_type(htreeitem, new ItemViewport(htreeitem)));
	return htreeitem;
}
HTREEITEM CVerticalTree::InsertItem(LPCTSTR lpszItem, HTREEITEM hParent, HTREEITEM hInsertAfter)
{
	HTREEITEM htreeitem = CTreeCtrl::InsertItem(lpszItem, hParent, hInsertAfter);
	//I presume that CTreeCtrl::InsertItem won't create same handles,
	//if the key is in use this will lead to memory leaks...
	m_HTree2VP.insert(HTREE2VPMAP::value_type(htreeitem, new ItemViewport(htreeitem)));
	return htreeitem;
}
HTREEITEM CVerticalTree::InsertItem(LPCTSTR lpszItem, int nImage, int nSelectedImage, HTREEITEM hParent, HTREEITEM hInsertAfter)
{
	HTREEITEM htreeitem = CTreeCtrl::InsertItem(lpszItem, nImage, nSelectedImage, hParent, hInsertAfter);
	//I presume that CTreeCtrl::InsertItem won't create same handles,
	//if the key is in use this will lead to memory leaks...
	m_HTree2VP.insert(HTREE2VPMAP::value_type(htreeitem, new ItemViewport(htreeitem)));
	return htreeitem;
}
BOOL CVerticalTree::DeleteItem(HTREEITEM hItem)
{
	BOOL ret = CTreeCtrl::DeleteItem(hItem);
	HTREE2VPMAP::iterator it = m_HTree2VP.find(hItem);
	if (it != m_HTree2VP.end())
	{
		delete (*it).second;
		m_HTree2VP.erase(it);
	}
	return ret;
}
BOOL CVerticalTree::DeleteAllItems()
{
	DeleteAllViewportItems();
	return CTreeCtrl::DeleteAllItems();
}

void CVerticalTree::DeleteAllViewportItems()
{
	HTREE2VPMAP::iterator it = m_HTree2VP.begin();
	while (it != m_HTree2VP.end())
	{
		delete (*it).second;
		it++;
	}

}

__forceinline CVerticalTree::ItemViewport* CVerticalTree::GetViewport(HTREEITEM hItem)
{
	HTREE2VPMAP::iterator it = m_HTree2VP.find(hItem);
	if (it != m_HTree2VP.end())
		return (*it).second;
	return NULL;
}

void CVerticalTree::ResetScrollbars()
{
	CRect rect;
	GetClientRect(rect);

	if(rect.Height() > m_YMax + 8 )
	{
		ShowScrollBar( SB_VERT, FALSE );	// Hide it
		SetScrollPos( SB_VERT, 0 );
	}
	else
	{
		SCROLLINFO	si;
		si.cbSize = sizeof(SCROLLINFO);
		si.fMask = SIF_PAGE | SIF_RANGE;
		si.nPage = rect.Height();
		si.nMax = m_YMax + 8;
		si.nMin = 0 ;

		SetScrollInfo( SB_VERT, &si );
		EnableScrollBarCtrl( SB_VERT, TRUE );
	}


	if(rect.Width() > m_XMax + 8 )
	{
		ShowScrollBar( SB_HORZ, FALSE );	// Hide it
		SetScrollPos( SB_HORZ, 0 );
	}
	else
	{
		SCROLLINFO	si;
		si.cbSize = sizeof(SCROLLINFO);
		si.fMask = SIF_PAGE | SIF_RANGE;
		si.nPage = rect.Width();
		si.nMax = m_XMax + 8;
		si.nMin = 0 ;

		SetScrollInfo( SB_HORZ, &si );
		EnableScrollBarCtrl( SB_HORZ, TRUE );
	}
}

void CVerticalTree::OnSizing(UINT fwSide, LPRECT pRect) 
{
	if (!m_bCustomDrawn)
		CTreeCtrl::OnSizing(fwSide, pRect);
}

void CVerticalTree::OnSize(UINT nType, int cx, int cy) 
{
	if (!m_bCustomDrawn)
		CTreeCtrl::OnSize(nType, cx, cy);
	else
	{
		ResetScrollbars();
		Invalidate();
	}	
}