Custom Draw TreeView in VB.NET

Public Class TreeViewEx
	Inherits TreeView

	Private Structure RECT
		Friend left As Integer
		Friend top As Integer
		Friend right As Integer
		Friend bottom As Integer
	End Structure

	Private Structure NMHDR
		Friend hwndFrom As IntPtr
		Friend idFrom As IntPtr
		Friend code As Integer
	End Structure

	Private Structure NMCUSTOMDRAW
		Friend hdr As NMHDR
		Friend dwDrawStage As Integer
		Friend hdc As IntPtr
		Friend rc As RECT
		Friend dwItemSpec As IntPtr
		Friend uItemState As Integer
		Friend lItemlParam As IntPtr
	End Structure

	Private Structure NMTVCUSTOMDRAW
		Friend nmcd As NMCUSTOMDRAW
		Friend clrText As Integer
		Friend clrTextBk As Integer
		Friend iLevel As Integer
	End Structure

	Private Function HandleNotify(ByVal m As System.Windows.Forms.Message) As Integer

		' Reference:
		' http://msdn.microsoft.com/library/default.asp?url=/library/en-us/shellcc/platform/commctls/custdraw/custdraw.asp

		Const NM_FIRST As Integer = 0
		Const NM_CUSTOMDRAW As Integer = (NM_FIRST - 12)

		' Drawstage flags
		Const CDDS_PREPAINT As Integer = &H1
		Const CDDS_POSTPAINT As Integer = &H2
		Const CDDS_PREERASE As Integer = &H3
		Const CDDS_POSTERASE As Integer = &H4

		Const CDDS_ITEM As Integer = &H10000
		Const CDDS_ITEMPREPAINT As Integer = (CDDS_ITEM Or CDDS_PREPAINT)
		Const CDDS_ITEMPOSTPAINT As Integer = (CDDS_ITEM Or CDDS_POSTPAINT)
		Const CDDS_ITEMPREERASE As Integer = (CDDS_ITEM Or CDDS_PREERASE)
		Const CDDS_ITEMPOSTERASE As Integer = (CDDS_ITEM Or CDDS_POSTERASE)
		Const CDDS_SUBITEM As Integer = &H20000

		' Custom draw return flags
		Const CDRF_DODEFAULT As Integer = &H0
		Const CDRF_NEWFONT As Integer = &H2
		Const CDRF_SKIPDEFAULT As Integer = &H4
		Const CDRF_NOTIFYPOSTPAINT As Integer = &H10
		Const CDRF_NOTIFYITEMDRAW As Integer = &H20
		Const CDRF_NOTIFYSUBITEMDRAW As Integer = &H20		 ' Flags are the same, we can distinguish by context
		Const CDRF_NOTIFYPOSTERASE As Integer = &H40

		Dim tNMHDR As NMHDR
		Dim tNMTVCUSTOMDRAW As NMTVCUSTOMDRAW
		Dim iResult As Integer = 0
		Dim objObject As Object
		Dim objGraphics As Graphics
		Dim objTreeNode As TreeNode
		Dim objTreeNodeEx As TreeNodeEx


		Try

			If Not m.LParam.Equals(IntPtr.Zero) Then

				objObject = m.GetLParam(GetType(NMHDR))

				If TypeOf objObject Is NMHDR Then

					tNMHDR = DirectCast(objObject, NMHDR)

					If tNMHDR.code = NM_CUSTOMDRAW Then

						objObject = m.GetLParam(GetType(NMTVCUSTOMDRAW))

						If TypeOf objObject Is NMTVCUSTOMDRAW Then

							tNMTVCUSTOMDRAW = DirectCast(objObject, NMTVCUSTOMDRAW)

							Select Case tNMTVCUSTOMDRAW.nmcd.dwDrawStage

								Case CDDS_PREPAINT
									iResult = CDRF_NOTIFYITEMDRAW

								Case CDDS_ITEMPREPAINT
									iResult = CDRF_NOTIFYPOSTPAINT

								Case CDDS_ITEMPOSTPAINT

									objTreeNode = TreeNode.FromHandle(Me, tNMTVCUSTOMDRAW.nmcd.dwItemSpec)
									objTreeNodeEx = DirectCast(objTreeNode, TreeNodeEx)

									objGraphics = Graphics.FromHdc(tNMTVCUSTOMDRAW.nmcd.hdc)

									PaintTreeNode(objTreeNodeEx, objGraphics)

									objGraphics.Dispose()

									iResult = CDRF_DODEFAULT

							End Select

						End If

					End If

				End If

			End If

		Catch ex As Exception
			System.Windows.Forms.MessageBox.Show(ex.ToString)
		End Try

		Return iResult

	End Function

	Private Sub PaintTreeNode(ByVal objTreeNodeEx As TreeNodeEx, ByVal objGraphics As Graphics)

		Dim objBoldFont As Font
		Dim objNormalFont As Font
		Dim objBackgroundBrush As Brush
		Dim objForegroundBrush As Brush
		Dim objPen As Pen
		Dim sngPosX As Single
		Dim sngPosY As Single
		Dim sngWidth As Single
		Dim sngHeight As Single
		Dim sTextPortion1 As String
		Dim sTextPortion2 As String
		Dim sTextPortion3 As String
		Dim iTreeNodeWidth As Integer
		Dim sngTreeNodeWidthPortion1 As Single
		Dim sngTreeNodeWidthPortion2 As Single
		Dim sngTreeNodeWidthPortion3 As Single

		' Get the fonts
		objNormalFont = Me.Font
		objBoldFont = New Font(objNormalFont, FontStyle.Bold)

		' Get the position and dimensions of the treenode
		sngWidth = objTreeNodeEx.Bounds.Width - 2
		sngHeight = objTreeNodeEx.Bounds.Height - 2
		sngPosX = objTreeNodeEx.Bounds.X + 1
		sngPosY = objTreeNodeEx.Bounds.Y + 1

		' If there is text in bold, we need to get the width of each portion
		If objTreeNodeEx.BoldTextInitialPosition >= 0 And objTreeNodeEx.BoldTextLength > 0 Then

			' Get the text before the bold portion
			sTextPortion1 = objTreeNodeEx.Text.Substring(0, objTreeNodeEx.BoldTextInitialPosition)

			' Get the text of the bold portion
			sTextPortion2 = objTreeNodeEx.Text.Substring(objTreeNodeEx.BoldTextInitialPosition, objTreeNodeEx.BoldTextLength)

			' Get the text after the bold portion
			If objTreeNodeEx.BoldTextInitialPosition + objTreeNodeEx.BoldTextLength < objTreeNodeEx.Text.Length Then
				sTextPortion3 = objTreeNodeEx.Text.Substring(objTreeNodeEx.BoldTextInitialPosition + objTreeNodeEx.BoldTextLength)
			End If

			' Get the width of each portion, taking into account the font
			sngTreeNodeWidthPortion1 = MeasureCorrectedTextWidth(objGraphics, objNormalFont, sngWidth, sngHeight, sTextPortion1)
			sngTreeNodeWidthPortion2 = MeasureCorrectedTextWidth(objGraphics, objBoldFont, sngWidth, sngHeight, sTextPortion2)
			sngTreeNodeWidthPortion3 = MeasureCorrectedTextWidth(objGraphics, objNormalFont, sngWidth, sngHeight, sTextPortion3)

			' Get the total width
			iTreeNodeWidth = CType(sngTreeNodeWidthPortion1 + sngTreeNodeWidthPortion2 + sngTreeNodeWidthPortion3, Integer)

		Else
			iTreeNodeWidth = objTreeNodeEx.Bounds.Width
		End If

		' Make a correction to ensure always a correct width
		iTreeNodeWidth += 6

		' Get the brushes. Note: we should take into account the BackColor and ForeColor of the treenode (left as exercise)
		If Me.SelectedNode Is objTreeNodeEx Then
			objBackgroundBrush = SystemBrushes.Highlight
			objForegroundBrush = SystemBrushes.HighlightText
		Else
			objBackgroundBrush = SystemBrushes.Window
			objForegroundBrush = SystemBrushes.WindowText
		End If

		' Fill the background rectangle
		objGraphics.FillRectangle(objBackgroundBrush, objTreeNodeEx.Bounds.X, objTreeNodeEx.Bounds.Y, iTreeNodeWidth, objTreeNodeEx.Bounds.Height)

		' Draw focus rectangle if it is the selected treenode
		If Me.SelectedNode Is objTreeNodeEx Then
			objPen = New Pen(Color.Gray, 1)
			objPen.DashStyle = Drawing.Drawing2D.DashStyle.Dot
			objGraphics.DrawRectangle(objPen, objTreeNodeEx.Bounds.X, objTreeNodeEx.Bounds.Y, iTreeNodeWidth - 1, objTreeNodeEx.Bounds.Height - 1)
			objPen.Dispose()
		End If

		' Draw the text
		If objTreeNodeEx.BoldTextInitialPosition >= 0 Then

			' Part 1
			If sTextPortion1 <> "" Then
				sngPosX += PaintText(objGraphics, sTextPortion1, objNormalFont, objForegroundBrush, sngPosX, sngPosY, sngWidth, sngHeight)
			End If

			' Part 2
			If sTextPortion2 <> "" Then
				sngPosX += PaintText(objGraphics, sTextPortion2, objBoldFont, objForegroundBrush, sngPosX, sngPosY, sngWidth, sngHeight)
			End If

			' Part 3
			If sTextPortion3 <> "" Then

				' If the first character after the bold portion is a space character, add an extra pixel
				If sTextPortion3.StartsWith(" ") Then
					sngPosX += 1
				End If

				sngPosX += PaintText(objGraphics, sTextPortion3, objNormalFont, objForegroundBrush, sngPosX, sngPosY, sngWidth, sngHeight)

			End If

		Else
			sngPosX += PaintText(objGraphics, objTreeNodeEx.Text, objNormalFont, objForegroundBrush, sngPosX, sngPosY, sngWidth, sngHeight)
		End If

		objBoldFont.Dispose()

	End Sub

	Private Function PaintText(ByVal objGraphics As Graphics, ByVal sText As String, ByVal objFont As Font, ByVal objBrush As Brush, ByVal sngPosX As Single, ByVal sngPosY As Single, ByVal sngWidth As Single, ByVal sngHeight As Single) As Single

		Dim sngNewWidth As Single

		objGraphics.DrawString(sText, objFont, objBrush, sngPosX, sngPosY)

		sngNewWidth = MeasureCorrectedTextWidth(objGraphics, objFont, sngWidth, sngHeight, sText)

		Return sngNewWidth

	End Function

	Private Function MeasureCorrectedTextWidth(ByVal objGraphics As Graphics, ByVal objFont As Font, ByVal sngWidth As Single, ByVal sngHeight As Single, ByVal sText As String) As Single

		Dim sngSingleTextWidth As Single
		Dim sngDoubleTextWidth As Single
		Dim sngResult As Single = 0.0

		If sText <> "" Then

			' The measurement routine (MeasureCharacterRanges) adds some extra pixels to the result, that we want to discard.
			' To do this, we meausure the string and the string duplicated, and the difference is the measure that we want.
			' That is:
			' A = X + C
			' B = 2X + C
			' Where A and B are known (the measures) and C is unknown. We are interested in X, which is X = B - A
			sngSingleTextWidth = MeasureTextWidth(objGraphics, objFont, sngWidth, sngHeight, sText)
			sngDoubleTextWidth = MeasureTextWidth(objGraphics, objFont, sngWidth * 2, sngHeight, sText & sText)

			sngResult = sngDoubleTextWidth - sngSingleTextWidth

		End If

		Return sngResult

	End Function

	Private Function MeasureTextWidth(ByVal objGraphics As Graphics, ByVal objFont As Font, ByVal sngWidth As Single, ByVal sngHeight As Single, ByVal sText As String) As Single

		Dim sngResult As Single
		Dim colCharacterRanges(0) As CharacterRange
		Dim colRegions(1) As Region
		Dim objStringFormat As StringFormat
		Dim objLayoutRectangleF As RectangleF
		Dim objMeasureRectangleF As RectangleF

		objStringFormat = New StringFormat()

		' Allow enough width for the bold case
		If objFont.Bold Then
			sngWidth = sngWidth * 2
		End If

		objLayoutRectangleF = New RectangleF(0, 0, sngWidth, sngHeight)

		colCharacterRanges(0) = New CharacterRange(0, sText.Length)
		objStringFormat.SetMeasurableCharacterRanges(colCharacterRanges)
		colRegions = objGraphics.MeasureCharacterRanges(sText, objFont, objLayoutRectangleF, objStringFormat)
		objMeasureRectangleF = colRegions(0).GetBounds(objGraphics)
		sngResult = objMeasureRectangleF.Width

		Return sngResult

	End Function

	Protected Overrides Sub WndProc(ByRef m As System.Windows.Forms.Message)

		Const WM_NOTIFY As Integer = &H4E

		Dim iResult As Integer
		Dim bHandled As Boolean = False

		If m.Msg = (&H2000 Or WM_NOTIFY) Then		 ' It is the reflected WM_NOTIFY message sent to the parent

			If m.WParam.Equals(Me.Handle) Then
				iResult = HandleNotify(m)
				m.Result = New IntPtr(iResult)
				bHandled = True
			End If

		End If

		If Not bHandled Then
			MyBase.WndProc(m)
		End If

	End Sub

End Class