From de9bb7e6d8ef0133ee19ce46985e140a863b163f Mon Sep 17 00:00:00 2001
From: Jack Christensen <jack@jackchristensen.com>
Date: Sat, 29 Apr 2017 11:01:54 -0500
Subject: [PATCH] Use flyweight pattern for pgproto3 messages

---
 frontend.go | 103 +++++++++++++++++++++++++++++++++++++---------------
 1 file changed, 73 insertions(+), 30 deletions(-)

diff --git a/frontend.go b/frontend.go
index c1dec461..df67b718 100644
--- a/frontend.go
+++ b/frontend.go
@@ -12,6 +12,29 @@ import (
 type Frontend struct {
 	cr *chunkreader.ChunkReader
 	w  io.Writer
+
+	// Backend message flyweights
+	authentication       Authentication
+	backendKeyData       BackendKeyData
+	bindComplete         BindComplete
+	closeComplete        CloseComplete
+	commandComplete      CommandComplete
+	copyBothResponse     CopyBothResponse
+	copyData             CopyData
+	copyInResponse       CopyInResponse
+	copyOutResponse      CopyOutResponse
+	dataRow              DataRow
+	emptyQueryResponse   EmptyQueryResponse
+	errorResponse        ErrorResponse
+	functionCallResponse FunctionCallResponse
+	noData               NoData
+	noticeResponse       NoticeResponse
+	notificationResponse NotificationResponse
+	parameterDescription ParameterDescription
+	parameterStatus      ParameterStatus
+	parseComplete        ParseComplete
+	readyForQuery        ReadyForQuery
+	rowDescription       RowDescription
 }
 
 func NewFrontend(r io.Reader, w io.Writer) (*Frontend, error) {
@@ -24,30 +47,6 @@ func (b *Frontend) Send(msg FrontendMessage) error {
 }
 
 func (b *Frontend) Receive() (BackendMessage, error) {
-	backendMessages := map[byte]BackendMessage{
-		'1': &ParseComplete{},
-		'2': &BindComplete{},
-		'3': &CloseComplete{},
-		'A': &NotificationResponse{},
-		'C': &CommandComplete{},
-		'd': &CopyData{},
-		'D': &DataRow{},
-		'E': &ErrorResponse{},
-		'G': &CopyInResponse{},
-		'H': &CopyOutResponse{},
-		'I': &EmptyQueryResponse{},
-		'K': &BackendKeyData{},
-		'n': &NoData{},
-		'N': &NoticeResponse{},
-		'R': &Authentication{},
-		'S': &ParameterStatus{},
-		't': &ParameterDescription{},
-		'T': &RowDescription{},
-		'V': &FunctionCallResponse{},
-		'W': &CopyBothResponse{},
-		'Z': &ReadyForQuery{},
-	}
-
 	header, err := b.cr.Next(5)
 	if err != nil {
 		return nil, err
@@ -56,15 +55,59 @@ func (b *Frontend) Receive() (BackendMessage, error) {
 	msgType := header[0]
 	bodyLen := int(binary.BigEndian.Uint32(header[1:])) - 4
 
+	var msg BackendMessage
+	switch msgType {
+	case '1':
+		msg = &b.parseComplete
+	case '2':
+		msg = &b.bindComplete
+	case '3':
+		msg = &b.closeComplete
+	case 'A':
+		msg = &b.notificationResponse
+	case 'C':
+		msg = &b.commandComplete
+	case 'd':
+		msg = &b.copyData
+	case 'D':
+		msg = &b.dataRow
+	case 'E':
+		msg = &b.errorResponse
+	case 'G':
+		msg = &b.copyInResponse
+	case 'H':
+		msg = &b.copyOutResponse
+	case 'I':
+		msg = &b.emptyQueryResponse
+	case 'K':
+		msg = &b.backendKeyData
+	case 'n':
+		msg = &b.noData
+	case 'N':
+		msg = &b.noticeResponse
+	case 'R':
+		msg = &b.authentication
+	case 'S':
+		msg = &b.parameterStatus
+	case 't':
+		msg = &b.parameterDescription
+	case 'T':
+		msg = &b.rowDescription
+	case 'V':
+		msg = &b.functionCallResponse
+	case 'W':
+		msg = &b.copyBothResponse
+	case 'Z':
+		msg = &b.readyForQuery
+	default:
+		return nil, fmt.Errorf("unknown message type: %c", msgType)
+	}
+
 	msgBody, err := b.cr.Next(bodyLen)
 	if err != nil {
 		return nil, err
 	}
 
-	if msg, ok := backendMessages[msgType]; ok {
-		err = msg.UnmarshalBinary(msgBody)
-		return msg, err
-	}
-
-	return nil, fmt.Errorf("unknown message type: %c", msgType)
+	err = msg.UnmarshalBinary(msgBody)
+	return msg, err
 }