4 files changed, 36 insertions(+), 39 deletions(-)

M lib/dhall/ast.rb
M lib/dhall/binary.rb
M lib/dhall/typecheck.rb
M test/test_typechecker.rb

@@ 360,20 360,29 @@ module Dhall
 
 	class Optional < Expression
 		include(ValueSemantics.for_attributes do
-			value Expression
-			type Either(nil, Expression), default: nil
+			value      Expression
+			value_type Either(nil, Expression), default: nil
 		end)
 
 		def self.for(value, type: nil)
 			if value.nil?
-				OptionalNone.new(type: type)
+				OptionalNone.new(value_type: value_type)
 			else
-				Optional.new(value: value, type: type)
+				Optional.new(value: value, value_type: type)
 			end
 		end
 
+		def type
+			return unless value_type
+
+			Dhall::Application.new(
+				function: Dhall::Variable["Optional"],
+				argument: value_type
+			)
+		end
+
 		def map(type: nil, &block)
-			with(value: block[value], type: type)
+			with(value: block[value], value_type: value_type)
 		end
 
 		def reduce(_, &block)


@@ 381,17 390,17 @@ module Dhall
 		end
 
 		def as_json
-			[5, type&.as_json, value.as_json]
+			[5, value_type&.as_json, value.as_json]
 		end
 	end
 
 	class OptionalNone < Optional
 		include(ValueSemantics.for_attributes do
-			type Expression
+			value_type Expression
 		end)
 
 		def map(type: nil)
-			type.nil? ? self : with(type: type)
+			type.nil? ? self : with(value_type: value_type)
 		end
 
 		def reduce(z)


@@ 399,7 408,7 @@ module Dhall
 		end
 
 		def as_json
-			[0, Variable["None"].as_json, type.as_json]
+			[0, Variable["None"].as_json, value_type.as_json]
 		end
 	end
 

@@ 90,11 90,11 @@ module Dhall
 	class Optional
 		def self.decode(type, value=nil)
 			if value.nil?
-				OptionalNone.new(type: Dhall.decode(type))
+				OptionalNone.new(value_type: Dhall.decode(type))
 			else
 				Optional.new(
-					value: Dhall.decode(value),
-					type:  type.nil? ? type : Dhall.decode(type)
+					value:       Dhall.decode(value),
+					value_type:  type.nil? ? type : Dhall.decode(type)
 				)
 			end
 		end

@@ 436,18 436,13 @@ module Dhall
 			end
 
 			def annotate(context)
-				type_type = TypeChecker.for(@expr.type).annotate(context).type
-				if type_type != Dhall::Variable["Type"]
-					raise TypeError, "OptionalNone element type not of type Type"
-				end
-
-				Dhall::TypeAnnotation.new(
-					value: @expr,
-					type:  Dhall::Application.new(
-						function: Dhall::Variable["Optional"],
-						argument: @expr.type
-					)
+				TypeChecker.assert(
+					TypeChecker.for(@expr.value_type).annotate(context).type,
+					Dhall::Variable["Type"],
+					"OptionalNone element type not of type Type"
 				)
+
+				@expr
 			end
 		end
 


@@ 457,23 452,16 @@ module Dhall
 			end
 
 			def annotate(context)
-				asome = @some.map(type: @some.type) do |el|
+				asome = @some.map(type: @some.value_type) do |el|
 					TypeChecker.for(el).annotate(context)
 				end
-				some = asome.with(type: asome.value.type)
+				some = asome.with(value_type: asome.value.type)
 
-				type_type = TypeChecker.for(some.type).annotate(context).type
-				if type_type != Dhall::Variable["Type"]
-					raise TypeError, "Some type no of type Type, was: #{type_type}"
-				end
+				type_type = TypeChecker.for(some.value_type).annotate(context).type
+				TypeChecker.assert type_type, Dhall::Variable["Type"],
+				                   "Some type not of type Type, was: #{type_type}"
 
-				Dhall::TypeAnnotation.new(
-					value: some,
-					type:  Dhall::Application.new(
-						function: Dhall::Variable["Optional"],
-						argument: some.type
-					)
-				)
+				some
 			end
 		end
 

@@ 15,14 15,14 @@ class TestTypechecker < Minitest::Test
 		next if test =~ /prelude/
 
 		define_method("test_#{test}") do
-			assert_kind_of(
-				Dhall::TypeAnnotation,
+			assert_respond_to(
 				Dhall::TypeChecker.for(
 					Dhall::TypeAnnotation.new(
 						value: Dhall.from_binary(path.binread),
 						type:  Dhall.from_binary((TESTS + "#{test}B.dhallb").binread)
 					)
-				).annotate(Dhall::TypeChecker::Context.new)
+				).annotate(Dhall::TypeChecker::Context.new),
+				:type
 			)
 		end
 	end