From 7ee3c5bda0ee8bd6b13b17739d43a3cae5c9b392 Mon Sep 17 00:00:00 2001
From: 0x19 <nevio@eiger.co>
Date: Thu, 17 Aug 2023 21:17:22 +0200
Subject: [PATCH] Ups missing functionality for sources that are requred for
 compiler to work

---
 sources.go      | 19 +++++++++++++++++++
 sources_test.go |  9 +++++++++
 2 files changed, 28 insertions(+)

diff --git a/sources.go b/sources.go
index a5e29086..168e3bb0 100644
--- a/sources.go
+++ b/sources.go
@@ -327,6 +327,25 @@ func (s *Sources) TruncateDir(path string) error {
 	return nil
 }
 
+// GetSolidityVersion extracts the Solidity version from the entry source unit.
+func (s *Sources) GetSolidityVersion() (string, error) {
+	// Get the entry source unit by its name
+	entrySourceUnit := s.GetSourceUnitByName(s.EntrySourceUnitName)
+	if entrySourceUnit == nil {
+		return "", fmt.Errorf("entry source unit not found")
+	}
+
+	// Use a regular expression to match the pragma solidity statement
+	// This regex will match versions like ^0.x.x and extract only 0.x.x
+	re := regexp.MustCompile(`pragma solidity\s*\^?(\d+\.\d+\.\d+);`)
+	match := re.FindStringSubmatch(entrySourceUnit.Content)
+	if len(match) < 2 {
+		return "", fmt.Errorf("solidity version not found in entry source unit")
+	}
+
+	return match[1], nil
+}
+
 // handleImports extracts import statements from the source unit and adds them to the sources.
 func (s *Sources) handleImports(sourceUnit *SourceUnit) ([]*SourceUnit, error) {
 	imports := extractImports(sourceUnit.Content)
diff --git a/sources_test.go b/sources_test.go
index 80e76332..f1d4615f 100644
--- a/sources_test.go
+++ b/sources_test.go
@@ -16,6 +16,7 @@ func TestSources(t *testing.T) {
 		sources       *Sources
 		expected      string
 		expectedUnits int
+		noSourceUnit  bool
 	}{
 		{
 			name: "Test Case 1",
@@ -37,6 +38,7 @@ func TestSources(t *testing.T) {
 			},
 			expected:      "Content of Source 1\n\nContent of Source 2",
 			expectedUnits: 2,
+			noSourceUnit:  true,
 		},
 		{
 			name: "Openzeppelin import",
@@ -53,6 +55,7 @@ func TestSources(t *testing.T) {
 			},
 			expected:      tests.ReadContractFileForTestFromRootPath(t, "contracts/cheelee/Combined").Content,
 			expectedUnits: 15,
+			noSourceUnit:  true,
 		},
 		{
 			name: "OpenZeppelin ERC20 Test",
@@ -115,6 +118,12 @@ func TestSources(t *testing.T) {
 			assert.True(t, testCase.sources.SourceUnitExists(testCase.sources.SourceUnits[0].Name))
 			assert.NotNil(t, testCase.sources.GetSourceUnitByName(testCase.sources.SourceUnits[0].Name))
 			assert.NotNil(t, testCase.sources.GetSourceUnitByPath(testCase.sources.SourceUnits[0].Path))
+
+			if !testCase.noSourceUnit {
+				version, err := testCase.sources.GetSolidityVersion()
+				assert.NoError(t, err)
+				assert.NotEmpty(t, version)
+			}
 		})
 	}
 }