Implement study related methods for BlackHoleStorage

storage_blackhole.go

@@ -3,6 +3,9 @@ package goptuna
 import (
 	"errors"
 	"sync"
+	"time"
+
+	"github.com/c-bata/goptuna"
 )
 
 var (
@@ -29,7 +32,7 @@ func NewBlackholeStorage(n int) *BlackHoleStorage {
 	}
 }
 
-// BlackholeStorage is an in-memory storage, but designed for over 10k+ trials.
+// BlackholeStorage is an in-memory storage, but designed for over 10k trials.
 // This storage just holds the latest N trials you specified.
 type BlackHoleStorage struct {
 	direction   StudyDirection
@@ -43,52 +46,155 @@ type BlackHoleStorage struct {
 	mu          sync.RWMutex
 }
 
+// CreateNewStudy creates study and returns studyID.
 func (s *BlackHoleStorage) CreateNewStudy(name string) (int, error) {
-	panic("implement me")
+	if name != "" {
+		s.studyName = name
+	}
+	return InMemoryStorageStudyID, nil
 }
 
+// DeleteStudy deletes a study.
 func (s *BlackHoleStorage) DeleteStudy(studyID int) error {
-	panic("implement me")
+	if !s.checkStudyID(studyID) {
+		return goptuna.ErrInvalidStudyID
+	}
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	s.direction = goptuna.StudyDirectionMinimize
+	s.trials = make([]goptuna.FrozenTrial, 0, 128)
+	s.userAttrs = make(map[string]string, 8)
+	s.systemAttrs = make(map[string]string, 8)
+	s.studyName = DefaultStudyNamePrefix + InMemoryStorageStudyUUID
+	return nil
 }
 
+// SetStudyDirection sets study direction of the objective.
 func (s *BlackHoleStorage) SetStudyDirection(studyID int, direction StudyDirection) error {
-	panic("implement me")
+	if !s.checkStudyID(studyID) {
+		return goptuna.ErrInvalidStudyID
+	}
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	s.direction = direction
+	return nil
 }
 
+// SetStudyUserAttr to store the value for the user.
 func (s *BlackHoleStorage) SetStudyUserAttr(studyID int, key string, value string) error {
-	panic("implement me")
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	s.userAttrs[key] = value
+	return nil
 }
 
+// SetStudySystemAttr to store the value for the system.
 func (s *BlackHoleStorage) SetStudySystemAttr(studyID int, key string, value string) error {
-	panic("implement me")
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	s.systemAttrs[key] = value
+	return nil
 }
 
+// GetStudyIDFromName return the study id from study name.
 func (s *BlackHoleStorage) GetStudyIDFromName(name string) (int, error) {
-	panic("implement me")
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	if name != s.studyName {
+		return -1, goptuna.ErrNotFound
+	}
+	return InMemoryStorageStudyID, nil
 }
 
+// GetStudyIDFromTrialID return the study id from trial id.
 func (s *BlackHoleStorage) GetStudyIDFromTrialID(trialID int) (int, error) {
-	panic("implement me")
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	// counter represents an id for next trial.
+	if s.counter > trialID {
+		return InMemoryStorageStudyID, nil
+	}
+	return -1, goptuna.ErrNotFound
 }
 
+// GetStudyNameFromID return the study name from study id.
 func (s *BlackHoleStorage) GetStudyNameFromID(studyID int) (string, error) {
-	panic("implement me")
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	if !s.checkStudyID(studyID) {
+		return "", goptuna.ErrNotFound
+	}
+	return s.studyName, nil
 }
 
+// GetStudyDirection returns study direction of the objective.
 func (s *BlackHoleStorage) GetStudyDirection(studyID int) (StudyDirection, error) {
-	panic("implement me")
+	if !s.checkStudyID(studyID) {
+		return goptuna.StudyDirectionMinimize, goptuna.ErrInvalidStudyID
+	}
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	return s.direction, nil
 }
 
+// GetStudyUserAttrs to restore the attributes for the user.
 func (s *BlackHoleStorage) GetStudyUserAttrs(studyID int) (map[string]string, error) {
-	panic("implement me")
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	n := make(map[string]string, len(s.userAttrs))
+	for k := range s.userAttrs {
+		n[k] = s.userAttrs[k]
+	}
+	return n, nil
 }
 
+// GetStudySystemAttrs to restore the attributes for the system.
 func (s *BlackHoleStorage) GetStudySystemAttrs(studyID int) (map[string]string, error) {
-	panic("implement me")
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	n := make(map[string]string, len(s.systemAttrs))
+	for k := range s.systemAttrs {
+		n[k] = s.systemAttrs[k]
+	}
+	return n, nil
 }
 
+// GetAllStudySummaries returns all study summaries.
 func (s *BlackHoleStorage) GetAllStudySummaries() ([]StudySummary, error) {
-	panic("implement me")
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	var datetimeStart time.Time
+	sa := make(map[string]string, len(s.systemAttrs))
+	for k := range s.systemAttrs {
+		sa[k] = s.systemAttrs[k]
+	}
+	ua := make(map[string]string, len(s.userAttrs))
+	for k := range s.userAttrs {
+		ua[k] = s.userAttrs[k]
+	}
+
+	return []StudySummary{
+		{
+			ID:            InMemoryStorageStudyID,
+			Name:          s.studyName,
+			Direction:     s.direction,
+			BestTrial:     s.bestTrial,
+			UserAttrs:     ua,
+			SystemAttrs:   sa,
+			DatetimeStart: datetimeStart,
+		},
+	}, nil
 }
 
 func (s *BlackHoleStorage) CreateNewTrial(studyID int) (int, error) {
@@ -155,3 +261,7 @@ func (s *BlackHoleStorage) GetTrialUserAttrs(trialID int) (map[string]string, er
 func (s *BlackHoleStorage) GetTrialSystemAttrs(trialID int) (map[string]string, error) {
 	panic("implement me")
 }
+
+func (s *BlackHoleStorage) checkStudyID(studyID int) bool {
+	return studyID == InMemoryStorageStudyID
+}