Browse Source

Added possibility to get all unitIds recursively

Dmitry Yu Okunev 5 years ago
parent
commit
4e5747c265
1 changed files with 25 additions and 1 deletions
  1. 25 1
      unit.go

+ 25 - 1
unit.go

@@ -102,6 +102,7 @@ func (unit Unit) IsChildrenReady() bool {
 }
 
 type Units []Unit
+type UnitPtrs []*Unit
 
 func (units Units) ToPersNumberMap() map[int]*Unit {
 	unitMap := map[int]*Unit{}
@@ -172,6 +173,10 @@ func (unit *Unit) DoRecursive(f func(*Unit, interface{}) bool, arg interface{})
 }
 
 func (units Units) DoRecursive(f func(*Unit, interface{}) bool, arg interface{}) bool {
+	return units.ToPtrSlice().DoRecursive(f, arg)
+}
+
+func (units UnitPtrs) DoRecursive(f func(*Unit, interface{}) bool, arg interface{}) bool {
 	for _, unit := range units {
 		unit.DoRecursive(f, arg)
 	}
@@ -179,6 +184,21 @@ func (units Units) DoRecursive(f func(*Unit, interface{}) bool, arg interface{})
 	return true
 }
 
+func (units Units) ToPtrSlice() (result UnitPtrs) {
+	for idx, _ := range units {
+		result = append(result, &units[idx])
+	}
+	return
+}
+
+func (units Units) GetRecursive() (result UnitPtrs) {
+	units.DoRecursive(func(unit *Unit, arg interface{}) bool {
+		result = append(result, unit)
+		return true
+	}, nil)
+	return
+}
+
 func (u Unit) GetChildrenPtrs() []*Unit {
 	if !u.childrenReady {
 		panic("unit is not ready for method GetChildrenPtrs: it's required to call CalculateTree method, first")
@@ -391,7 +411,11 @@ func (units Units) GetPersNumbers() (persNumbers []int) {
 	return
 }
 
-func (units Units) GetUnitIds() (unitIds []int) {
+func (units Units) GetUnitIds() []int {
+	return units.ToPtrSlice().GetUnitIds()
+}
+
+func (units UnitPtrs) GetUnitIds() (unitIds []int) {
 	for _, unit := range units {
 		unitIds = append(unitIds, unit.Id)
 	}